MAINT: annotate out as NDArray, remove scattered isinstance checks

Author: ev-br

torch_np/_binary_ufuncs.py

@@ -1,6 +1,8 @@
 from . import _helpers
 from ._detail import _binary_ufuncs
-from ._normalizations import ArrayLike, DTypeLike, SubokLike, normalizer
+from ._normalizations import ArrayLike, DTypeLike, SubokLike, NDArray, normalizer
+from typing import Optional
+
 
 __all__ = [
     name for name in dir(_binary_ufuncs) if not name.startswith("_") and name != "torch"
@@ -18,7 +20,7 @@ def wrapped(
         x1: ArrayLike,
         x2: ArrayLike,
         /,
-        out=None,
+        out: Optional[NDArray] = None,
         *,
         where=True,
         casting="same_kind",

torch_np/_funcs.py

@@ -9,10 +9,12 @@
     ArrayLike,
     AxisLike,
     DTypeLike,
+    NDArray,
     SubokLike,
     UnpackedSeqArrayLike,
     normalizer,
 )
+from typing import Optional
 
 
 @normalizer
@@ -32,7 +34,7 @@ def clip(
     a: ArrayLike,
     min: Optional[ArrayLike] = None,
     max: Optional[ArrayLike] = None,
-    out=None,
+    out: Optional[NDArray] = None,
 ):
     # np.clip requires both a_min and a_max not None, while ndarray.clip allows
     # one of them to be None. Follow the more lax version.
@@ -57,7 +59,7 @@ def diagonal(a: ArrayLike, offset=0, axis1=0, axis2=1):
 
 
 @normalizer
-def trace(a: ArrayLike, offset=0, axis1=0, axis2=1, dtype: DTypeLike = None, out=None):
+def trace(a: ArrayLike, offset=0, axis1=0, axis2=1, dtype: DTypeLike = None, out: Optional[NDArray] = None):
     result = _impl.trace(a, offset, axis1, axis2, dtype)
     return _helpers.result_or_out(result, out)
 
@@ -112,7 +114,7 @@ def vdot(a: ArrayLike, b: ArrayLike, /):
 
 
 @normalizer
-def dot(a: ArrayLike, b: ArrayLike, out=None):
+def dot(a: ArrayLike, b: ArrayLike, out: Optional[NDArray] = None):
     result = _impl.dot(a, b)
     return _helpers.result_or_out(result, out)
 
@@ -211,7 +213,7 @@ def imag(a: ArrayLike):
 
 
 @normalizer
-def round_(a: ArrayLike, decimals=0, out=None):
+def round_(a: ArrayLike, decimals=0, out: Optional[NDArray]=None):
     result = _impl.round(a, decimals)
     return _helpers.result_or_out(result, out)
 
@@ -231,7 +233,7 @@ def sum(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out=None,
+    out: Optional[NDArray]=None,
     keepdims=NoValue,
     initial=NoValue,
     where=NoValue,
@@ -247,7 +249,7 @@ def prod(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out=None,
+    out: Optional[NDArray]=None,
     keepdims=NoValue,
     initial=NoValue,
     where=NoValue,
@@ -266,7 +268,7 @@ def mean(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out=None,
+    out: Optional[NDArray]=None,
     keepdims=NoValue,
     *,
     where=NoValue,
@@ -282,7 +284,7 @@ def var(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out=None,
+    out: Optional[NDArray]=None,
     ddof=0,
     keepdims=NoValue,
     *,
@@ -299,7 +301,7 @@ def std(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out=None,
+    out: Optional[NDArray]=None,
     ddof=0,
     keepdims=NoValue,
     *,
@@ -312,13 +314,13 @@ def std(
 
 
 @normalizer
-def argmin(a: ArrayLike, axis: AxisLike = None, out=None, *, keepdims=NoValue):
+def argmin(a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = None, *, keepdims=NoValue):
     result = _reductions.argmin(a, axis=axis, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
 @normalizer
-def argmax(a: ArrayLike, axis: AxisLike = None, out=None, *, keepdims=NoValue):
+def argmax(a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = None, *, keepdims=NoValue):
     result = _reductions.argmax(a, axis=axis, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
@@ -327,7 +329,7 @@ def argmax(a: ArrayLike, axis: AxisLike = None, out=None, *, keepdims=NoValue):
 def amax(
     a: ArrayLike,
     axis: AxisLike = None,
-    out=None,
+    out: Optional[NDArray] = None,
     keepdims=NoValue,
     initial=NoValue,
     where=NoValue,
@@ -345,7 +347,7 @@ def amax(
 def amin(
     a: ArrayLike,
     axis: AxisLike = None,
-    out=None,
+    out: Optional[NDArray] = None,
     keepdims=NoValue,
     initial=NoValue,
     where=NoValue,
@@ -360,22 +362,22 @@ def amin(
 
 
 @normalizer
-def ptp(a: ArrayLike, axis: AxisLike = None, out=None, keepdims=NoValue):
+def ptp(a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = None, keepdims=NoValue):
     result = _reductions.ptp(a, axis=axis, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
 @normalizer
 def all(
-    a: ArrayLike, axis: AxisLike = None, out=None, keepdims=NoValue, *, where=NoValue
+    a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = None, keepdims=NoValue, *, where=NoValue
 ):
     result = _reductions.all(a, axis=axis, where=where, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
 @normalizer
 def any(
-    a: ArrayLike, axis: AxisLike = None, out=None, keepdims=NoValue, *, where=NoValue
+    a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = None, keepdims=NoValue, *, where=NoValue
 ):
     result = _reductions.any(a, axis=axis, where=where, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
@@ -388,13 +390,13 @@ def count_nonzero(a: ArrayLike, axis: AxisLike = None, *, keepdims=False):
 
 
 @normalizer
-def cumsum(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out=None):
+def cumsum(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out: Optional[NDArray] = None):
     result = _reductions.cumsum(a, axis=axis, dtype=dtype)
     return _helpers.result_or_out(result, out)
 
 
 @normalizer
-def cumprod(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out=None):
+def cumprod(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out: Optional[NDArray] = None):
     result = _reductions.cumprod(a, axis=axis, dtype=dtype)
     return _helpers.result_or_out(result, out)
 
@@ -407,7 +409,7 @@ def quantile(
     a: ArrayLike,
     q: ArrayLike,
     axis: AxisLike = None,
-    out=None,
+    out: Optional[NDArray] = None,
     overwrite_input=False,
     method="linear",
     keepdims=False,

torch_np/_helpers.py

@@ -29,11 +29,6 @@ def cast_and_broadcast(tensors, out, casting):
     if out is None:
         return tensors
     else:
-        from ._ndarray import asarray, ndarray
-
-        if not isinstance(out, ndarray):
-            raise TypeError("Return arrays must be of ArrayType")
-
         tensors = _util.cast_and_broadcast(
             tensors, out.dtype.type.torch_dtype, out.shape, casting
         )
@@ -72,11 +67,7 @@ def result_or_out(result_tensor, out_array=None, promote_scalar=False):
             result_tensor is placed into the out array.
     This weirdness is used e.g. in `np.percentile`
     """
-    from ._ndarray import asarray, ndarray
-
     if out_array is not None:
-        if not isinstance(out_array, ndarray):
-            raise TypeError("Return arrays must be of ArrayType")
         if result_tensor.shape != out_array.shape:
             can_fit = result_tensor.numel() == 1 and out_array.ndim == 0
             if promote_scalar and can_fit:
@@ -90,7 +81,7 @@ def result_or_out(result_tensor, out_array=None, promote_scalar=False):
         out_tensor.copy_(result_tensor)
         return out_array
     else:
-        return asarray(result_tensor)
+        return array_from(result_tensor)
 
 
 def array_from(tensor, base=None):

torch_np/_normalizations.py

@@ -12,6 +12,7 @@
 DTypeLike = typing.TypeVar("DTypeLike")
 SubokLike = typing.TypeVar("SubokLike")
 AxisLike = typing.TypeVar("AxisLike")
+NDArray = typing.TypeVar("NDarray")
 
 # annotate e.g. atleast_1d(*arys)
 UnpackedSeqArrayLike = typing.TypeVar("UnpackedSeqArrayLike")
@@ -60,11 +61,24 @@ def normalize_axis_like(arg, name=None):
     return arg
 
 
+def normalize_ndarray(arg, name=None):
+    if arg is None:
+        return arg
+
+    from ._ndarray import ndarray
+
+    if not isinstance(arg, ndarray):
+        raise TypeError("'out' must be an array")
+    return arg
+
+
+
 normalizers = {
     ArrayLike: normalize_array_like,
     Optional[ArrayLike]: normalize_optional_array_like,
     Sequence[ArrayLike]: normalize_seq_array_like,
     UnpackedSeqArrayLike: normalize_seq_array_like,  # cf handling in normalize
+    Optional[NDArray]: normalize_ndarray,
     DTypeLike: normalize_dtype,
     SubokLike: normalize_subok_like,
     AxisLike: normalize_axis_like,

torch_np/_unary_ufuncs.py

@@ -4,7 +4,8 @@
 
 from . import _helpers
 from ._detail import _unary_ufuncs
-from ._normalizations import ArrayLike, DTypeLike, SubokLike, normalizer
+from ._normalizations import ArrayLike, DTypeLike, SubokLike, NDArray, normalizer
+from typing import Optional
 
 __all__ = [
     name for name in dir(_unary_ufuncs) if not name.startswith("_") and name != "torch"
@@ -21,7 +22,7 @@ def deco_unary_ufunc(torch_func):
     def wrapped(
         x: ArrayLike,
         /,
-        out=None,
+        out: Optional[NDArray] = None,
         *,
         where=True,
         casting="same_kind",

torch_np/_wrapper.py

@@ -13,14 +13,15 @@
 from ._detail import implementations as _impl
 from ._ndarray import array, asarray, maybe_set_base, ndarray
 
-### XXX: order the imports DAG
 from ._normalizations import (
     ArrayLike,
     DTypeLike,
     SubokLike,
     UnpackedSeqArrayLike,
+    NDArray,
     normalizer,
 )
+from typing import Optional
 
 from . import _dtypes, _helpers, _decorators  # isort: skip  # XXX
 
@@ -108,9 +109,6 @@ def _concat_check(tup, dtype, out):
         raise ValueError("need at least one array to concatenate")
 
     if out is not None:
-        if not isinstance(out, ndarray):
-            raise ValueError("'out' must be an array")
-
         if dtype is not None:
             # mimic numpy
             raise TypeError(
@@ -123,7 +121,7 @@ def _concat_check(tup, dtype, out):
 def concatenate(
     ar_tuple: Sequence[ArrayLike],
     axis=0,
-    out=None,
+    out: Optional[NDArray]=None,
     dtype: DTypeLike = None,
     casting="same_kind",
 ):
@@ -173,7 +171,7 @@ def column_stack(
 def stack(
     arrays: Sequence[ArrayLike],
     axis=0,
-    out=None,
+    out:  Optional[NDArray] = None,
     *,
     dtype: DTypeLike = None,
     casting="same_kind",
@@ -666,7 +664,7 @@ def percentile(
     a,
     q,
     axis=None,
-    out=None,
+    out:  Optional[NDArray] = None,
     overwrite_input=False,
     method="linear",
     keepdims=False,
@@ -678,7 +676,7 @@ def percentile(
     )
 
 
-def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
+def median(a, axis=None, out:  Optional[NDArray] = None, overwrite_input=False, keepdims=False):
     return _funcs.quantile(
         a, 0.5, axis=axis, overwrite_input=overwrite_input, out=out, keepdims=keepdims
     )
@@ -691,7 +689,7 @@ def inner(a: ArrayLike, b: ArrayLike, /):
 
 
 @normalizer
-def outer(a: ArrayLike, b: ArrayLike, out=None):
+def outer(a: ArrayLike, b: ArrayLike, out:  Optional[NDArray] = None):
     result = torch.outer(a, b)
     return _helpers.result_or_out(result, out)
 
@@ -704,7 +702,7 @@ def nanmean(
     a: ArrayLike,
     axis=None,
     dtype: DTypeLike = None,
-    out=None,
+    out:  Optional[NDArray] = None,
     keepdims=NoValue,
     *,
     where=NoValue,
@@ -847,13 +845,13 @@ def isrealobj(x: ArrayLike):
 
 
 @normalizer
-def isneginf(x: ArrayLike, out=None):
+def isneginf(x: ArrayLike, out:  Optional[NDArray] = None):
     result = torch.isneginf(x, out=out)
     return _helpers.array_from(result)
 
 
 @normalizer
-def isposinf(x: ArrayLike, out=None):
+def isposinf(x: ArrayLike, out:  Optional[NDArray] = None):
     result = torch.isposinf(x, out=out)
     return _helpers.array_from(result)