Merge wrapper and implementations

torch_np/__init__.py

@@ -1,13 +1,11 @@
 from . import random
 from ._binary_ufuncs import *
-from ._detail._index_tricks import *
 from ._detail._util import AxisError, UFuncTypeError
 from ._dtypes import *
 from ._funcs import *
 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

@@ -86,7 +86,7 @@ def divmod(
     out1: Optional[NDArray] = None,
     out2: Optional[NDArray] = None,
     /,
-    out: Optional[tuple[NDArray]] = (None, None),
+    out: tuple[Optional[NDArray], Optional[NDArray]] = (None, None),
     *,
     where=True,
     casting="same_kind",

torch_np/_detail/__init__.py

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

torch_np/_detail/_util.py

@@ -170,36 +170,6 @@ def typecast_tensors(tensors, target_dtype, casting):
     return tuple(cast_tensors)
 
 
-def axis_expand_func(func, tensor, axis, *args, **kwds):
-    """Generically handle axis arguments in reductions."""
-    if axis is not None:
-        if not isinstance(axis, (list, tuple)):
-            axis = (axis,)
-        axis = normalize_axis_tuple(axis, tensor.ndim)
-
-    if axis == ():
-        newshape = expand_shape(tensor.shape, axis=0)
-        tensor = tensor.reshape(newshape)
-        axis = (0,)
-
-    result = func(tensor, axis=axis, *args, **kwds)
-
-    return result
-
-
-def axis_ravel_func(func, tensor, axis, *args, **kwds):
-    """Generically handle axis arguments in cumsum/cumprod."""
-    if axis is not None:
-        axis = normalize_axis_index(axis, tensor.ndim)
-
-    tensors, axis = axis_none_ravel(tensor, axis=axis)
-    tensor = tensors[0]
-
-    result = func(tensor, axis=axis, *args, **kwds)
-
-    return result
-
-
 def _coerce_to_tensor(obj, dtype=None, copy=False, ndmin=0):
     """The core logic of the array(...) function.