Typing: More informative return typing for ops

[jahall]

Description

The return type of .graph.op.Op.__call__ is Union[Variable, List[Variable]]. However, some child classes (e.g. .tensor.basic.Alloc) only ever return a single output, in which case an output type of Variable would be more intuitive (aside from the case where return_list=True) as the union causes annoying issues with mypy as per seen in this PR.

I feel like something should be possible with typing.overload but can't really see how to implement.

However, for the case that return_list=True I think we could do the following:

from typing import Literal, overload

# -- snip --

    @overload
    def __call__(self, *inputs: Any, return_list: Literal[True] = True, **kwargs) -> List[Variable]:
        ...
    
    def __call__(self, *inputs: Any, **kwargs) -> Union[Variable, List[Variable]]:
        r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.

# -- snip --

[ricardoV94]

That's a small use case but we can definitely add it. Do you want to open a PR?