MAINT: make _detail a package, hide individual implementation files and import from _detail

Author: ev-br

torch_np/_detail/__init__.py

@@ -0,0 +1,6 @@
+from ._flips import *
+from ._reductions import *
+
+# leading underscore (ndarray.flatten yes, np.flatten no)
+from .implementations import *
+from .implementations import _flatten

torch_np/_detail/_reductions.py

@@ -10,7 +10,7 @@
 
 from . import _dtypes_impl, _util
 
-NoValue = None
+NoValue = _util.NoValue
 
 
 import functools

torch_np/_detail/_util.py

@@ -7,6 +7,7 @@
 
 from . import _dtypes_impl
 
+NoValue = None
 
 # https://github.com/numpy/numpy/blob/v1.23.0/numpy/distutils/misc_util.py#L497-L504
 def is_sequence(seq):

torch_np/_funcs.py

@@ -2,9 +2,9 @@
 
 import torch
 
+from . import _detail as _impl
 from . import _helpers
-from ._detail import _flips, _reductions, _util
-from ._detail import implementations as _impl
+from ._detail import _util
 from ._normalizations import (
     ArrayLike,
     AxisLike,
@@ -15,6 +15,8 @@
     normalizer,
 )
 
+NoValue = _util.NoValue
+
 
 @normalizer
 def nonzero(a: ArrayLike):
@@ -159,13 +161,13 @@ def moveaxis(a: ArrayLike, source, destination):
 
 @normalizer
 def swapaxes(a: ArrayLike, axis1, axis2):
-    result = _flips.swapaxes(a, axis1, axis2)
+    result = _impl.swapaxes(a, axis1, axis2)
     return _helpers.array_from(result)
 
 
 @normalizer
 def rollaxis(a: ArrayLike, axis, start=0):
-    result = _flips.rollaxis(a, axis, start)
+    result = _impl.rollaxis(a, axis, start)
     return _helpers.array_from(result)
 
 
@@ -231,9 +233,6 @@ def round_(a: ArrayLike, decimals=0, out: Optional[NDArray] = None):
 # ### reductions ###
 
 
-NoValue = None  # FIXME
-
-
 @normalizer
 def sum(
     a: ArrayLike,
@@ -244,7 +243,7 @@ def sum(
     initial=NoValue,
     where=NoValue,
 ):
-    result = _reductions.sum(
+    result = _impl.sum(
         a, axis=axis, dtype=dtype, initial=initial, where=where, keepdims=keepdims
     )
     return _helpers.result_or_out(result, out)
@@ -260,7 +259,7 @@ def prod(
     initial=NoValue,
     where=NoValue,
 ):
-    result = _reductions.prod(
+    result = _impl.prod(
         a, axis=axis, dtype=dtype, initial=initial, where=where, keepdims=keepdims
     )
     return _helpers.result_or_out(result, out)
@@ -279,9 +278,7 @@ def mean(
     *,
     where=NoValue,
 ):
-    result = _reductions.mean(
-        a, axis=axis, dtype=dtype, where=NoValue, keepdims=keepdims
-    )
+    result = _impl.mean(a, axis=axis, dtype=dtype, where=NoValue, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
@@ -296,7 +293,7 @@ def var(
     *,
     where=NoValue,
 ):
-    result = _reductions.var(
+    result = _impl.var(
         a, axis=axis, dtype=dtype, ddof=ddof, where=where, keepdims=keepdims
     )
     return _helpers.result_or_out(result, out)
@@ -313,7 +310,7 @@ def std(
     *,
     where=NoValue,
 ):
-    result = _reductions.std(
+    result = _impl.std(
         a, axis=axis, dtype=dtype, ddof=ddof, where=where, keepdims=keepdims
     )
     return _helpers.result_or_out(result, out)
@@ -327,7 +324,7 @@ def argmin(
     *,
     keepdims=NoValue,
 ):
-    result = _reductions.argmin(a, axis=axis, keepdims=keepdims)
+    result = _impl.argmin(a, axis=axis, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
@@ -339,7 +336,7 @@ def argmax(
     *,
     keepdims=NoValue,
 ):
-    result = _reductions.argmax(a, axis=axis, keepdims=keepdims)
+    result = _impl.argmax(a, axis=axis, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
@@ -352,9 +349,7 @@ def amax(
     initial=NoValue,
     where=NoValue,
 ):
-    result = _reductions.max(
-        a, axis=axis, initial=initial, where=where, keepdims=keepdims
-    )
+    result = _impl.max(a, axis=axis, initial=initial, where=where, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
@@ -370,9 +365,7 @@ def amin(
     initial=NoValue,
     where=NoValue,
 ):
-    result = _reductions.min(
-        a, axis=axis, initial=initial, where=where, keepdims=keepdims
-    )
+    result = _impl.min(a, axis=axis, initial=initial, where=where, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
@@ -383,7 +376,7 @@ def amin(
 def ptp(
     a: ArrayLike, axis: AxisLike = None, out: Optional[NDArray] = None, keepdims=NoValue
 ):
-    result = _reductions.ptp(a, axis=axis, keepdims=keepdims)
+    result = _impl.ptp(a, axis=axis, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
@@ -396,7 +389,7 @@ def all(
     *,
     where=NoValue,
 ):
-    result = _reductions.all(a, axis=axis, where=where, keepdims=keepdims)
+    result = _impl.all(a, axis=axis, where=where, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
@@ -409,13 +402,13 @@ def any(
     *,
     where=NoValue,
 ):
-    result = _reductions.any(a, axis=axis, where=where, keepdims=keepdims)
+    result = _impl.any(a, axis=axis, where=where, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
 
 @normalizer
 def count_nonzero(a: ArrayLike, axis: AxisLike = None, *, keepdims=False):
-    result = _reductions.count_nonzero(a, axis=axis, keepdims=keepdims)
+    result = _impl.count_nonzero(a, axis=axis, keepdims=keepdims)
     return _helpers.array_from(result)
 
 
@@ -426,7 +419,7 @@ def cumsum(
     dtype: DTypeLike = None,
     out: Optional[NDArray] = None,
 ):
-    result = _reductions.cumsum(a, axis=axis, dtype=dtype)
+    result = _impl.cumsum(a, axis=axis, dtype=dtype)
     return _helpers.result_or_out(result, out)
 
 
@@ -437,7 +430,7 @@ def cumprod(
     dtype: DTypeLike = None,
     out: Optional[NDArray] = None,
 ):
-    result = _reductions.cumprod(a, axis=axis, dtype=dtype)
+    result = _impl.cumprod(a, axis=axis, dtype=dtype)
     return _helpers.result_or_out(result, out)
 
 
@@ -459,5 +452,5 @@ def quantile(
     if interpolation is not None:
         raise ValueError("'interpolation' argument is deprecated; use 'method' instead")
 
-    result = _reductions.quantile(a, q, axis, method=method, keepdims=keepdims)
+    result = _impl.quantile(a, q, axis, method=method, keepdims=keepdims)
     return _helpers.result_or_out(result, out, promote_scalar=True)

torch_np/_ndarray.py

@@ -3,7 +3,7 @@
 import torch
 
 from . import _binary_ufuncs, _dtypes, _funcs, _helpers, _unary_ufuncs
-from ._detail import _dtypes_impl, _flips, _reductions, _util
+from ._detail import _dtypes_impl, _util
 from ._detail import implementations as _impl
 
 newaxis = None

torch_np/_wrapper.py

@@ -8,9 +8,10 @@
 
 import torch
 
-from . import _funcs, _helpers
-from ._detail import _dtypes_impl, _flips, _reductions, _util
-from ._detail import implementations as _impl
+from . import _decorators
+from . import _detail as _impl
+from . import _dtypes, _funcs, _helpers
+from ._detail import _dtypes_impl, _util
 from ._ndarray import asarray
 from ._normalizations import (
     ArrayLike,
@@ -21,6 +22,8 @@
     normalizer,
 )
 
+NoValue = _util.NoValue
+
 # Things to decide on (punt for now)
 #
 # 1. Q: What are the return types of wrapper functions: plain torch.Tensors or
@@ -56,9 +59,6 @@
 #   - optional out arg
 
 
-NoValue = None
-
-
 ###### array creation routines
 
 
@@ -499,25 +499,25 @@ def expand_dims(a: ArrayLike, axis):
 
 @normalizer
 def flip(m: ArrayLike, axis=None):
-    result = _flips.flip(m, axis)
+    result = _impl.flip(m, axis)
     return _helpers.array_from(result)
 
 
 @normalizer
 def flipud(m: ArrayLike):
-    result = _flips.flipud(m)
+    result = _impl.flipud(m)
     return _helpers.array_from(result)
 
 
 @normalizer
 def fliplr(m: ArrayLike):
-    result = _flips.fliplr(m)
+    result = _impl.fliplr(m)
     return _helpers.array_from(result)
 
 
 @normalizer
 def rot90(m: ArrayLike, k=1, axes=(0, 1)):
-    result = _flips.rot90(m, k, axes)
+    result = _impl.rot90(m, k, axes)
     return _helpers.array_from(result)
 
 
@@ -640,9 +640,7 @@ def average(
     *,
     keepdims=NoValue,
 ):
-    result, wsum = _reductions.average(
-        a, axis, weights, returned=returned, keepdims=keepdims
-    )
+    result, wsum = _impl.average(a, axis, weights, returned=returned, keepdims=keepdims)
     if returned:
         return _helpers.tuple_arrays_from((result, wsum))
     else: