wrap returns out annotation

[ev-br]

Use a return annotation for functions with an out= argument. Other returns are wrapped generically.

Detect the out argument in signatures and implement the out= semantics.

supersedes and closes gh-83

[lezcano]

I really don't like having to return both the result and the out. It feels very unnatural and counterintuitive to me, as, to implement the function, you have to effectively change its signature (e.g. in C++ you'd go from returning a Tensor to returning a std::tuple<Tensor, Tensor&>?!). Even more, the wrap_tensors function now does too many things: It wraps result and it implements the out behaviour. These are two independent things.

I believe a better way to do this would be:

1. Make functions always return just the result and never annotate the result i.e. leave things how they were before.
2. Leave wrap_tensor as it was.
3. Rewrite this bit as

def maybe_copy_to(out, result, promote_scalar_result=False):
    if out is None:
        return result
    elif isinstance(result, torch.Tensor):
        if result.shape != out.shape:
            can_fit = result.numel() == 1 and out.ndim == 0
            if promote_scalar_result and can_fit:
                result = result.squeeze()
            else:
                raise ValueError(
                    f"Bad size of the out array: out.shape = {out.shape}"
                    f" while result.shape = {result.shape}."
                )
        out.copy_(result)
        return out
    elif isinstance(result, (tuple, list)):
        return type(result)(map(maybe_copy_to, zip(result, out)))
    else:
        assert False  # We should never hit this path

...
if "out" in params:
  out = sig.bind(*args, **kwargs).arguments.get("out")
  result = maybe_copy_out(result, out, promote_scalar_result)
result = wrap_tensors(result)

[lezcano]

Also, the annotation for out doesn't seem correct, as it may be a List[NDArray] or Tuple[NDArray] as well, right?

[ev-br]

I'm going to prepare a version without return annotations.

[ev-br]

Force pushed a version without type annotations, almost identical to #93 (comment).
I was wanting to avoid dispatching on argument names ("out"), but the alternatives seem clunky, so let's roll with this and see if want to improve at leisure.

One possibly semi-missing bit is divmod (again!), which has several ways of accepting out arrays. Working out all details of what numpy really does requires taking over numpy/core/tests/test_umath.py. This is rather large and rather hairy, so best not conflate it with the cleanup that is this PR.

[lezcano]

I agree that dispatching on name is not the best. Another way of doing this would be to have an annotation Output = Union[NDArray, Tuple[NDArray...], Tuple[Any...]] (for example, it's not like we're using them for typecheking really) and then dispatch on that. Either works for me.

About divmod... my my. Yeah, your one-off solution LGTM.

[lezcano]

nit tuple[Optional[NDArray], Optional[NDArray]].

[ev-br]

Why? It's either (None, None) or (NDArray, NDArray); can't have (out1, None) (I think)

[lezcano]

tuple, as opposed, to list takes the number of arguments that it accepts explicitly. As such, there you are declaring there a tuple of length 1 and one NDArray.
Then, Optional[tuple[NDArray, NDArray]] accepts either a tuple of ndarrays (fine) or None (not (None, None)).
What you want is to have a (partial) function that accepts tuple[Optional[NDArray], Optional[NDArray]] and errors out when just one of the elements of the tuple is None.

[lezcano]

Feel free to patch this small correction in any other PR.

[ev-br]

Ouch. python typing meta-language is seriously weird.

[lezcano]

Note that this is not specific to Python, the same would be true in C++. This would be an std::tuple<std::optional<array, std::optional<array>> rather than a std::optional<std::tuple<array, array>> :)

[lezcano]

As discussed ufunc_preprocess should not take an out parameter as setting an out= parameter should not affect the type of the computation. Or is this fixed in some other PR and we'll just need to update this one afterwards?

[ev-br]

It's still in the TODO pipeline. Will do after the current "stack" is done.

[ev-br]

Thanks Mario