MAINT: remove isort:skip directives (circ imports are well hidden now)

Author: ev-br

torch_np/__init__.py

@@ -1,4 +1,3 @@
-from ._wrapper import *  # isort: skip  # XXX: currently this prevents circular imports
 from . import random
 from ._binary_ufuncs import *
 from ._detail._index_tricks import *
@@ -8,6 +7,7 @@
 from ._getlimits import errstate, finfo, iinfo
 from ._ndarray import array, asarray, can_cast, ndarray, newaxis, result_type
 from ._unary_ufuncs import *
+from ._wrapper import *
 
 # from . import testing
 

torch_np/_binary_ufuncs.py

@@ -1,8 +1,8 @@
-from . import _helpers
-from ._detail import _binary_ufuncs
-from ._normalizations import ArrayLike, DTypeLike, SubokLike, NDArray, normalizer
 from typing import Optional
 
+from . import _helpers
+from ._detail import _binary_ufuncs
+from ._normalizations import ArrayLike, DTypeLike, NDArray, SubokLike, normalizer
 
 __all__ = [
     name for name in dir(_binary_ufuncs) if not name.startswith("_") and name != "torch"

torch_np/_detail/_reductions.py

@@ -4,6 +4,8 @@
 Anything here only deals with torch objects, e.g. "dtype" is a torch.dtype instance etc
 """
 
+import typing
+
 import torch
 
 from . import _dtypes_impl, _util
@@ -26,7 +28,9 @@ def wrapped(tensor, axis, *args, **kwds):
         if axis is not None:
             if not isinstance(axis, (list, tuple)):
                 if not isinstance(axis, typing.SupportsIndex):
-                    raise TypeError(f"{type(axis)=}, but should be a list/tuple or support operator.index()")
+                    raise TypeError(
+                        f"{type(axis)=}, but should be a list/tuple or support operator.index()"
+                    )
                 axis = (axis,)
             axis = _util.normalize_axis_tuple(axis, tensor.ndim)
 

torch_np/_funcs.py

@@ -14,7 +14,6 @@
     UnpackedSeqArrayLike,
     normalizer,
 )
-from typing import Optional
 
 
 @normalizer
@@ -59,7 +58,14 @@ 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: Optional[NDArray] = 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)
 
@@ -213,7 +219,7 @@ def imag(a: ArrayLike):
 
 
 @normalizer
-def round_(a: ArrayLike, decimals=0, out: Optional[NDArray]=None):
+def round_(a: ArrayLike, decimals=0, out: Optional[NDArray] = None):
     result = _impl.round(a, decimals)
     return _helpers.result_or_out(result, out)
 
@@ -233,7 +239,7 @@ def sum(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out: Optional[NDArray]=None,
+    out: Optional[NDArray] = None,
     keepdims=NoValue,
     initial=NoValue,
     where=NoValue,
@@ -249,7 +255,7 @@ def prod(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out: Optional[NDArray]=None,
+    out: Optional[NDArray] = None,
     keepdims=NoValue,
     initial=NoValue,
     where=NoValue,
@@ -268,7 +274,7 @@ def mean(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out: Optional[NDArray]=None,
+    out: Optional[NDArray] = None,
     keepdims=NoValue,
     *,
     where=NoValue,
@@ -284,7 +290,7 @@ def var(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out: Optional[NDArray]=None,
+    out: Optional[NDArray] = None,
     ddof=0,
     keepdims=NoValue,
     *,
@@ -301,7 +307,7 @@ def std(
     a: ArrayLike,
     axis: AxisLike = None,
     dtype: DTypeLike = None,
-    out: Optional[NDArray]=None,
+    out: Optional[NDArray] = None,
     ddof=0,
     keepdims=NoValue,
     *,
@@ -314,13 +320,25 @@ def std(
 
 
 @normalizer
-def argmin(a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = 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: Optional[NDArray] = 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)
 
@@ -362,22 +380,34 @@ def amin(
 
 
 @normalizer
-def ptp(a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = 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: Optional[NDArray] = 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: Optional[NDArray] = 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)
@@ -390,13 +420,23 @@ def count_nonzero(a: ArrayLike, axis: AxisLike = None, *, keepdims=False):
 
 
 @normalizer
-def cumsum(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out: Optional[NDArray] = 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: Optional[NDArray] = 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)
 

torch_np/_normalizations.py

@@ -72,7 +72,6 @@ def normalize_ndarray(arg, name=None):
     return arg
 
 
-
 normalizers = {
     ArrayLike: normalize_array_like,
     Optional[ArrayLike]: normalize_optional_array_like,

torch_np/_unary_ufuncs.py

@@ -2,10 +2,11 @@
 # from ._detail import _ufunc_impl
 
 
+from typing import Optional
+
 from . import _helpers
 from ._detail import _unary_ufuncs
-from ._normalizations import ArrayLike, DTypeLike, SubokLike, NDArray, normalizer
-from typing import Optional
+from ._normalizations import ArrayLike, DTypeLike, NDArray, SubokLike, normalizer
 
 __all__ = [
     name for name in dir(_unary_ufuncs) if not name.startswith("_") and name != "torch"

torch_np/_wrapper.py

@@ -8,23 +8,18 @@
 
 import torch
 
-from . import _funcs
+from . import _decorators, _dtypes, _funcs, _helpers
 from ._detail import _dtypes_impl, _flips, _reductions, _util
 from ._detail import implementations as _impl
 from ._ndarray import array, asarray, maybe_set_base, ndarray
-
 from ._normalizations import (
     ArrayLike,
     DTypeLike,
+    NDArray,
     SubokLike,
     UnpackedSeqArrayLike,
-    NDArray,
     normalizer,
 )
-from typing import Optional
-
-from . import _dtypes, _helpers, _decorators  # isort: skip  # XXX
-
 
 # Things to decide on (punt for now)
 #
@@ -121,7 +116,7 @@ def _concat_check(tup, dtype, out):
 def concatenate(
     ar_tuple: Sequence[ArrayLike],
     axis=0,
-    out: Optional[NDArray]=None,
+    out: Optional[NDArray] = None,
     dtype: DTypeLike = None,
     casting="same_kind",
 ):
@@ -171,7 +166,7 @@ def column_stack(
 def stack(
     arrays: Sequence[ArrayLike],
     axis=0,
-    out:  Optional[NDArray] = None,
+    out: Optional[NDArray] = None,
     *,
     dtype: DTypeLike = None,
     casting="same_kind",
@@ -664,7 +659,7 @@ def percentile(
     a,
     q,
     axis=None,
-    out:  Optional[NDArray] = None,
+    out: Optional[NDArray] = None,
     overwrite_input=False,
     method="linear",
     keepdims=False,
@@ -676,7 +671,9 @@ def percentile(
     )
 
 
-def median(a, axis=None, out:  Optional[NDArray] = 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
     )
@@ -689,7 +686,7 @@ def inner(a: ArrayLike, b: ArrayLike, /):
 
 
 @normalizer
-def outer(a: ArrayLike, b: ArrayLike, out:  Optional[NDArray] = None):
+def outer(a: ArrayLike, b: ArrayLike, out: Optional[NDArray] = None):
     result = torch.outer(a, b)
     return _helpers.result_or_out(result, out)
 
@@ -702,7 +699,7 @@ def nanmean(
     a: ArrayLike,
     axis=None,
     dtype: DTypeLike = None,
-    out:  Optional[NDArray] = None,
+    out: Optional[NDArray] = None,
     keepdims=NoValue,
     *,
     where=NoValue,
@@ -845,13 +842,13 @@ def isrealobj(x: ArrayLike):
 
 
 @normalizer
-def isneginf(x: ArrayLike, out:  Optional[NDArray] = 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:  Optional[NDArray] = None):
+def isposinf(x: ArrayLike, out: Optional[NDArray] = None):
     result = torch.isposinf(x, out=out)
     return _helpers.array_from(result)
 

torch_np/random.py

@@ -6,13 +6,13 @@
 
 """
 from math import sqrt
+from typing import Optional
 
 import torch
 
-from ._detail import _dtypes_impl, _util
 from . import _helpers
-from ._normalizations import normalizer, ArrayLike
-from typing import Optional
+from ._detail import _dtypes_impl, _util
+from ._normalizations import ArrayLike, normalizer
 
 _default_dtype = _dtypes_impl.default_float_dtype
 
@@ -96,7 +96,7 @@ def randint(low, high=None, size=None):
 
 
 @normalizer
-def choice(a: ArrayLike, size=None, replace=True, p: Optional[ArrayLike]=None):
+def choice(a: ArrayLike, size=None, replace=True, p: Optional[ArrayLike] = None):
 
     # https://stackoverflow.com/questions/59461811/random-choice-with-pytorch
     if a.numel() == 1: