MAINT: simplify/normalize average

Author: ev-br

torch_np/_detail/_reductions.py

@@ -292,7 +292,26 @@ def cumsum(tensor, axis, dtype=None):
     return result
 
 
-def average(a_tensor, axis, w_tensor):
+
+def average(a, axis, weights, returned=False, keepdims=False):
+    if weights is None:
+        result, wsum = average_noweights(a, axis, keepdims=keepdims)
+    else:
+        result, wsum = average_weights(a, axis, weights, keepdims=keepdims)
+
+    if returned:
+        if wsum.shape != result.shape:
+            wsum = torch.broadcast_to(wsum, result.shape).clone()
+    return result, wsum
+
+
+def average_noweights(a_tensor, axis, keepdims=False):
+    result = mean(a_tensor, axis=axis, keepdims=keepdims)
+    scl = torch.as_tensor(a_tensor.numel() / result.numel(), dtype=result.dtype)
+    return result, scl
+
+
+def average_weights(a_tensor, axis, w_tensor, keepdims=False):
 
     # dtype
     # FIXME: 1. use result_type
@@ -306,6 +325,9 @@ def average(a_tensor, axis, w_tensor):
     a_tensor = _util.cast_if_needed(a_tensor, result_dtype)
     w_tensor = _util.cast_if_needed(w_tensor, result_dtype)
 
+    # axis=None ravels, so store the originals to reuse with keepdims=True below
+    ax, ndim = axis, a_tensor.ndim
+
     # axis
     if axis is None:
         (a_tensor, w_tensor), axis = _util.axis_none_ravel(
@@ -334,6 +356,10 @@ def average(a_tensor, axis, w_tensor):
     denominator = w_tensor.sum(axis)
     result = numerator / denominator
 
+    # keepdims
+    if keepdims:
+        result = _util.apply_keepdims(result, ax, ndim)
+
     return result, denominator
 
 

torch_np/_funcs.py

@@ -461,6 +461,9 @@ def cumprod(a: ArrayLike, axis: AxisLike = None, dtype: DTypeLike = None, out=No
     return _helpers.result_or_out(result, out)
 
 
+cumproduct = cumprod
+
+
 @normalizer
 def quantile(
     a : ArrayLike,

torch_np/_wrapper.py

@@ -631,50 +631,16 @@ def tri(N, M=None, k=0, dtype: DTypeLike = float, *, like: SubokLike = None):
 
 ###### reductions
 
-
-def cumprod(a, axis=None, dtype=None, out=None):
-    arr = asarray(a)
-    return arr.cumprod(axis=axis, dtype=dtype, out=out)
-
-
-cumproduct = cumprod
-
-
-def cumsum(a, axis=None, dtype=None, out=None):
-    arr = asarray(a)
-    return arr.cumsum(axis=axis, dtype=dtype, out=out)
-
-
-def average(a, axis=None, weights=None, returned=False, *, keepdims=NoValue):
-
-    if weights is None:
-        result = _funcs.mean(a, axis=axis, keepdims=keepdims)
-        if returned:
-            scl = result.dtype.type(a.size / result.size)
-            return result, scl
-        return result
-
-    a_tensor, w_tensor = _helpers.to_tensors(a, weights)
-
-    result, wsum = _reductions.average(a_tensor, axis, w_tensor)
-
-    # keepdims
-    if keepdims:
-        result = _util.apply_keepdims(result, axis, a_tensor.ndim)
-
-    # returned
+@normalizer
+def average(a: ArrayLike, axis=None, weights: ArrayLike=None, returned=False, *, keepdims=NoValue):
+    result, wsum = _reductions.average(a, axis, weights, returned=returned, keepdims=keepdims)
     if returned:
-        scl = wsum
-        if scl.shape != result.shape:
-            scl = torch.broadcast_to(scl, result.shape).clone()
-
-        return _helpers.array_from(result), _helpers.array_from(scl)
-
+        return _helpers.tuple_arrays_from((result, wsum))
     else:
         return _helpers.array_from(result)
 
 
-# Normalizations (ArrayLike et al) are done in quantile.
+# Normalizations (ArrayLike et al) in percentile and median are done in `_funcs.py/quantile`.
 def percentile(
     a,
     q,
@@ -710,6 +676,8 @@ def outer(a: ArrayLike, b: ArrayLike, out=None):
     return _helpers.result_or_out(result, out)
 
 
+# ### FIXME: this is a stub
+
 @normalizer
 def nanmean(
     a: ArrayLike,