MAINT: cumsum/cumprod

Author: ev-br

torch_np/_detail/_reductions.py

@@ -48,6 +48,22 @@ def wrapped(tensor, axis=None, keepdims=NoValue, *args, **kwds):
         return result
     return wrapped
 
+
+def deco_axis_ravel(func):
+    """Generically handle 'axis=None ravels' behavior."""
+    @functools.wraps(func)
+    def wrapped(tensor, axis, *args, **kwds):
+        if axis is not None:
+            axis = _util.normalize_axis_index(axis, tensor.ndim)
+
+        tensors, axis = _util.axis_none_ravel(tensor, axis=axis)  # XXX: inline
+        tensor = tensors[0]
+
+        result = func(tensor, axis=axis, *args, **kwds)
+        return result
+    return wrapped
+
+
 ##################################3
 
 
@@ -252,6 +268,7 @@ def var(tensor, axis=None, dtype=None, ddof=0, *, where=NoValue):
 #   2. axis=None ravels (cf concatenate)
 
 
+@deco_axis_ravel
 def cumprod(tensor, axis, dtype=None):
     if dtype == torch.bool:
         dtype = _dtypes_impl.default_int_dtype
@@ -263,6 +280,7 @@ def cumprod(tensor, axis, dtype=None):
     return result
 
 
+@deco_axis_ravel
 def cumsum(tensor, axis, dtype=None):
     if dtype == torch.bool:
         dtype = _dtypes_impl.default_int_dtype

torch_np/_detail/_util.py

@@ -128,7 +128,7 @@ def apply_keepdims(tensor, axis, ndim):
 
 def axis_none_ravel(*tensors, axis=None):
     """Ravel the arrays if axis is none."""
-    # XXX: is only used at `concatenate`. Inline unless reused more widely
+    # XXX: is only used at `concatenate` and cumsum/cumprod. Inline unless reused more widely
     if axis is None:
         tensors = tuple(ar.ravel() for ar in tensors)
         return tensors, 0

torch_np/_funcs.py

@@ -448,3 +448,14 @@ def count_nonzero(a: ArrayLike, axis: AxisLike=None, *, keepdims=False):
     result = _reductions.count_nonzero(a, axis=axis, keepdims=keepdims)
     return _helpers.array_from(result)
 
+
+@normalizer
+def cumsum(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out=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):
+    result = _reductions.cumprod(a, axis=axis, dtype=dtype)
+    return _helpers.result_or_out(result, out)

torch_np/_ndarray.py

@@ -3,13 +3,6 @@
 import torch
 
 from . import _binary_ufuncs, _dtypes, _funcs, _helpers, _unary_ufuncs
-from ._decorators import (
-    NoValue,
-    axis_keepdims_wrapper,
-    axis_none_ravel_wrapper,
-    dtype_to_torch,
-    emulate_out_arg,
-)
 from ._detail import _dtypes_impl, _flips, _reductions, _util
 from ._detail import implementations as _impl
 
@@ -389,12 +382,8 @@ def sort(self, axis=-1, kind=None, order=None):
     var = _funcs.var
     std = _funcs.std
 
-    cumprod = emulate_out_arg(
-        axis_none_ravel_wrapper(dtype_to_torch(_reductions.cumprod))
-    )
-    cumsum = emulate_out_arg(
-        axis_none_ravel_wrapper(dtype_to_torch(_reductions.cumsum))
-    )
+    cumsum = _funcs.cumsum
+    cumprod = _funcs.cumprod
 
     ### indexing ###
     @staticmethod