MAINT: tuple_arrays_from or just an array

`_helpers.tuple_arrays_from` is now localized to the normalizer only

Author: ev-br

torch_np/_detail/implementations.py

@@ -372,7 +372,7 @@ def tile(tensor, reps):
 def atleast_1d(tensors):
     result = torch.atleast_1d(*tensors)
 
-    # match numpy: return a list not tuple; 
+    # 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)])

torch_np/_funcs.py

@@ -15,7 +15,6 @@
     normalizer,
 )
 
-
 NoValue = _util.NoValue
 
 

torch_np/_normalizations.py

@@ -168,7 +168,7 @@ def wrapped(*args, **kwds):
                 # 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
+                # a variadic return: 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))

torch_np/_wrapper.py

@@ -15,9 +15,9 @@
     ArrayLike,
     DTypeLike,
     NDArray,
+    NDArrayOrSequence,
     SubokLike,
     UnpackedSeqArrayLike,
-    NDArrayOrSequence,
     normalizer,
 )
 
@@ -522,13 +522,23 @@ def meshgrid(
     return output
 
 
-@normalizer
-def indices(dimensions, dtype: DTypeLike = int, sparse=False):
-    result = _impl.indices(dimensions, dtype=dtype, sparse=sparse)
+def indices(dimensions, dtype=int, sparse=False):
     if sparse:
-        return _helpers.tuple_arrays_from(result)
+        return _indices_sparse(dimensions, dtype)
     else:
-        return _helpers.array_from(result)
+        return _indices_full(dimensions, dtype)
+
+
+@normalizer
+def _indices_sparse(dimensions, dtype: DTypeLike = int) -> tuple[NDArray]:
+    result = _impl.indices(dimensions, dtype=dtype, sparse=True)
+    return result
+
+
+@normalizer
+def _indices_full(dimensions, dtype: DTypeLike = int) -> NDArray:
+    result = _impl.indices(dimensions, dtype=dtype, sparse=False)
+    return result
 
 
 @normalizer
@@ -581,20 +591,43 @@ def tri(N, M=None, k=0, dtype: DTypeLike = float, *, like: SubokLike = None) ->
 ###### reductions
 
 
+def average(a, axis=None, weights=None, returned=False, keepdims=NoValue):
+    if returned:
+        return _average_with_ret(
+            a, axis=axis, weights=weights, returned=True, keepdims=keepdims
+        )
+    else:
+        return _average_without_ret(
+            a, axis=axis, weights=weights, returned=False, keepdims=keepdims
+        )
+
+
+@normalizer
+def _average_with_ret(
+    a: ArrayLike,
+    axis=None,
+    weights: Optional[ArrayLike] = None,
+    returned=True,
+    *,
+    keepdims=NoValue,
+) -> tuple[NDArray]:
+    assert returned
+    result, wsum = _impl.average(a, axis, weights, returned=True, keepdims=keepdims)
+    return result, wsum
+
+
 @normalizer
-def average(
+def _average_without_ret(
     a: ArrayLike,
     axis=None,
-    weights: ArrayLike = None,
+    weights: Optional[ArrayLike] = None,
     returned=False,
     *,
     keepdims=NoValue,
-):
-    result, wsum = _impl.average(a, axis, weights, returned=returned, keepdims=keepdims)
-    if returned:
-        return _helpers.tuple_arrays_from((result, wsum))
-    else:
-        return _helpers.array_from(result)
+) -> NDArray:
+    assert not returned
+    result, wsum = _impl.average(a, axis, weights, returned=False, keepdims=keepdims)
+    return result
 
 
 @normalizer