MAINT: return annotations for tuples of NDArrays

Two cases
- just {tuple, list}[NDArrays]  (list is e.g. meshgrid)
- a variadic return of a single NDArray or a list/tuple (where, atleast_{1,2,3}d, unique)

Note that in the latter case the choice of the return type depends purely on
the number of input array_likes.

Author: ev-br

torch_np/_detail/implementations.py

@@ -369,6 +369,36 @@ def tile(tensor, reps):
     return result
 
 
+def atleast_1d(tensors):
+    result = torch.atleast_1d(*tensors)
+
+    # match numpy: return a list not tuple; 
+    # >>> np.atleast_2d(np.arange(3))
+    # array([[0, 1, 2]])          # a single 2D array
+    # >>> torch.atleast_2d([torch.arange(3)])
+    # (tensor([[0, 1, 2]]), )     # 1-element tuple of a 2D tensor
+    if isinstance(result, tuple):
+        return list(result)
+    else:
+        return result
+
+
+def atleast_2d(tensors):
+    result = torch.atleast_2d(*tensors)
+    if isinstance(result, tuple):
+        return list(result)
+    else:
+        return result
+
+
+def atleast_3d(tensors):
+    result = torch.atleast_3d(*tensors)
+    if isinstance(result, tuple):
+        return list(result)
+    else:
+        return result
+
+
 # #### cov & corrcoef
 
 

torch_np/_funcs.py

@@ -15,20 +15,22 @@
     normalizer,
 )
 
+
 NoValue = _util.NoValue
 
 
 @normalizer
-def nonzero(a: ArrayLike):
+def nonzero(a: ArrayLike) -> tuple[NDArray]:
     result = a.nonzero(as_tuple=True)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
 def argwhere(a: ArrayLike) -> NDArray:
     result = torch.argwhere(a)
     return result
 
+
 @normalizer
 def flatnonzero(a: ArrayLike) -> NDArray:
     result = a.ravel().nonzero(as_tuple=True)[0]
@@ -83,7 +85,9 @@ def trace(
 
 
 @normalizer
-def eye(N, M=None, k=0, dtype: DTypeLike = float, order="C", *, like: SubokLike = None) -> NDArray:
+def eye(
+    N, M=None, k=0, dtype: DTypeLike = float, order="C", *, like: SubokLike = None
+) -> NDArray:
     if order != "C":
         raise NotImplementedError
     result = _impl.eye(N, M, k, dtype)
@@ -108,15 +112,16 @@ def diagflat(v: ArrayLike, k=0)  -> NDArray:
     return result
 
 
-def diag_indices(n, ndim=2):
+@normalizer
+def diag_indices(n, ndim=2) -> tuple[NDArray]:
     result = _impl.diag_indices(n, ndim)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
-def diag_indices_from(arr: ArrayLike):
+def diag_indices_from(arr: ArrayLike) -> tuple[NDArray]:
     result = _impl.diag_indices_from(arr)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer

torch_np/_normalizations.py

@@ -2,7 +2,7 @@
 """
 import operator
 import typing
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Union
 
 import torch
 
@@ -18,6 +18,10 @@
 UnpackedSeqArrayLike = typing.TypeVar("UnpackedSeqArrayLike")
 
 
+# return value of atleast_1d et al: single array of a list/tuple of arrays
+NDArrayOrSequence = Union[NDArray, Sequence[NDArray]]
+
+
 import inspect
 
 from . import _dtypes
@@ -154,11 +158,23 @@ def wrapped(*args, **kwds):
             result = func(*ba.args, **ba.kwargs)
 
             # handle returns
-            return_annotation = sig.return_annotation
-            if return_annotation == NDArray:
+            r = sig.return_annotation
+            if r == NDArray:
                 return _helpers.array_from(result)
-            elif return_annotation ==  inspect._empty:
+            elif r == inspect._empty:
                 return result
+            elif hasattr(r, "__origin__") and r.__origin__ in (list, tuple):
+                # this is tuple[NDArray] or list[NDArray]
+                # XXX: change to separate tuple and list normalizers?
+                return r.__origin__(_helpers.tuple_arrays_from(result))
+            elif r == NDArrayOrSequence:
+                # this is a single NDArray or tuple/list of NDArrays, e.g. atleast_1d
+                if isinstance(result, (tuple, list)):
+                    seq = type(result)
+                    return seq(_helpers.tuple_arrays_from(result))
+                else:
+                    return _helpers.array_from(result)
+
             else:
                 raise ValueError(f"Unknown return annotation {return_annotation}")
 

torch_np/_wrapper.py

@@ -17,6 +17,7 @@
     NDArray,
     SubokLike,
     UnpackedSeqArrayLike,
+    NDArrayOrSequence,
     normalizer,
 )
 
@@ -69,30 +70,21 @@ def copy(a: ArrayLike, order="K", subok: SubokLike = False) -> NDArray:
 
 
 @normalizer
-def atleast_1d(*arys: UnpackedSeqArrayLike):
-    res = torch.atleast_1d(*arys)
-    if len(res) == 1:
-        return _helpers.array_from(res[0])
-    else:
-        return list(_helpers.tuple_arrays_from(res))
+def atleast_1d(*arys: UnpackedSeqArrayLike) -> NDArrayOrSequence:
+    result = _impl.atleast_1d(*arys)
+    return result
 
 
 @normalizer
-def atleast_2d(*arys: UnpackedSeqArrayLike):
-    res = torch.atleast_2d(*arys)
-    if len(res) == 1:
-        return _helpers.array_from(res[0])
-    else:
-        return list(_helpers.tuple_arrays_from(res))
+def atleast_2d(*arys: UnpackedSeqArrayLike) -> NDArrayOrSequence:
+    result = _impl.atleast_2d(*arys)
+    return result
 
 
 @normalizer
-def atleast_3d(*arys: UnpackedSeqArrayLike):
-    res = torch.atleast_3d(*arys)
-    if len(res) == 1:
-        return _helpers.array_from(res[0])
-    else:
-        return list(_helpers.tuple_arrays_from(res))
+def atleast_3d(*arys: UnpackedSeqArrayLike) -> NDArrayOrSequence:
+    result = _impl.atleast_3d(*arys)
+    return result
 
 
 def _concat_check(tup, dtype, out):
@@ -175,33 +167,33 @@ def stack(
 
 
 @normalizer
-def array_split(ary: ArrayLike, indices_or_sections, axis=0):
+def array_split(ary: ArrayLike, indices_or_sections, axis=0) -> tuple[NDArray]:
     result = _impl.split_helper(ary, indices_or_sections, axis)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
-def split(ary: ArrayLike, indices_or_sections, axis=0):
+def split(ary: ArrayLike, indices_or_sections, axis=0) -> tuple[NDArray]:
     result = _impl.split_helper(ary, indices_or_sections, axis, strict=True)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
-def hsplit(ary: ArrayLike, indices_or_sections):
+def hsplit(ary: ArrayLike, indices_or_sections) -> tuple[NDArray]:
     result = _impl.hsplit(ary, indices_or_sections)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
-def vsplit(ary: ArrayLike, indices_or_sections):
+def vsplit(ary: ArrayLike, indices_or_sections) -> tuple[NDArray]:
     result = _impl.vsplit(ary, indices_or_sections)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
-def dsplit(ary: ArrayLike, indices_or_sections):
+def dsplit(ary: ArrayLike, indices_or_sections) -> tuple[NDArray]:
     result = _impl.dsplit(ary, indices_or_sections)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
@@ -421,13 +413,9 @@ def where(
     x: Optional[ArrayLike] = None,
     y: Optional[ArrayLike] = None,
     /,
-):
+) -> NDArrayOrSequence:
     result = _impl.where(condition, x, y)
-    if isinstance(result, tuple):
-        # single-argument where(condition)
-        return _helpers.tuple_arrays_from(result)
-    else:
-        return _helpers.array_from(result)
+    return result
 
 
 ###### module-level queries of object properties
@@ -496,10 +484,12 @@ def broadcast_to(array: ArrayLike, shape, subok: SubokLike = False) -> NDArray:
 
 # YYY: pattern: tuple of arrays as input, tuple of arrays as output; cf nonzero
 @normalizer
-def broadcast_arrays(*args: UnpackedSeqArrayLike, subok: SubokLike = False):
+def broadcast_arrays(
+    *args: UnpackedSeqArrayLike, subok: SubokLike = False
+) -> tuple[NDArray]:
     args = args[0]  # undo the *args wrapping in normalizer
     res = torch.broadcast_tensors(*args)
-    return _helpers.tuple_arrays_from(res)
+    return res
 
 
 def unravel_index(indices, shape, order="C"):
@@ -524,11 +514,12 @@ def ravel_multi_index(multi_index, dims, mode="raise", order="C"):
 
 
 @normalizer
-def meshgrid(*xi: UnpackedSeqArrayLike, copy=True, sparse=False, indexing="xy"):
+def meshgrid(
+    *xi: UnpackedSeqArrayLike, copy=True, sparse=False, indexing="xy"
+) -> list[NDArray]:
     xi = xi[0]  # undo the *xi wrapping in normalizer
     output = _impl.meshgrid(*xi, copy=copy, sparse=sparse, indexing=indexing)
-    outp = _helpers.tuple_arrays_from(output)
-    return list(outp)  # match numpy, return a list
+    return output
 
 
 @normalizer
@@ -559,26 +550,26 @@ def triu(m: ArrayLike, k=0) -> NDArray:
     return result
 
 
-def tril_indices(n, k=0, m=None):
+def tril_indices(n, k=0, m=None) -> tuple[NDArray]:
     result = _impl.tril_indices(n, k, m)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
-def triu_indices(n, k=0, m=None):
+def triu_indices(n, k=0, m=None) -> tuple[NDArray]:
     result = _impl.triu_indices(n, k, m)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
-def tril_indices_from(arr: ArrayLike, k=0):
+def tril_indices_from(arr: ArrayLike, k=0) -> tuple[NDArray]:
     result = _impl.tril_indices_from(arr, k)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
-def triu_indices_from(arr: ArrayLike, k=0):
+def triu_indices_from(arr: ArrayLike, k=0) -> tuple[NDArray]:
     result = _impl.triu_indices_from(arr, k)
-    return _helpers.tuple_arrays_from(result)
+    return result
 
 
 @normalizer
@@ -859,7 +850,7 @@ def unique(
     axis=None,
     *,
     equal_nan=True,
-):
+) -> NDArrayOrSequence:
     result = _impl.unique(
         ar,
         return_index=return_index,
@@ -868,11 +859,7 @@ def unique(
         axis=axis,
         equal_nan=equal_nan,
     )
-
-    if isinstance(result, tuple):
-        return _helpers.tuple_arrays_from(result)
-    else:
-        return _helpers.array_from(result)
+    return result
 
 
 ###### mapping from numpy API objects to wrappers from this module ######