MAINT: add prod, var, std to follow sum

Author: ev-br

torch_np/_helpers.py

@@ -131,3 +131,13 @@ def to_tensors(*inputs):
     """Convert all ndarrays from `inputs` to tensors."""
     return tuple([value.get() if isinstance(value, ndarray) else value
             for value in inputs])
+
+
+def float_or_default(dtype, self_dtype):
+    """dtype helper for reductions."""
+    if dtype is None:
+        dtype = self_dtype
+    if _dtypes.is_integer(dtype):
+        dtype = _dtypes.default_float_type()
+    torch_dtype = _dtypes.torch_dtype_from(dtype)
+    return torch_dtype

torch_np/_ndarray.py

@@ -416,12 +416,7 @@ def mean(self, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoVal
         if where is not None:
             raise NotImplementedError
 
-        if dtype is None:
-            dtype = self.dtype
-        if _dtypes.is_integer(dtype):
-            dtype = _dtypes.default_float_type()
-        torch_dtype = _dtypes.torch_dtype_from(dtype)
-
+        torch_dtype = _helpers.float_or_default(dtype, self.dtype)
         if axis is None:
             result = self._tensor.mean(dtype=torch_dtype)
         else:
@@ -436,19 +431,57 @@ def sum(self, axis=None, dtype=None, out=None, keepdims=NoValue,
         if initial is not None or where is not None:
             raise NotImplementedError
 
-        if dtype is None:
-            dtype = self.dtype
-        if _dtypes.is_integer(dtype):
-            dtype = _dtypes.default_float_type()
-        torch_dtype = _dtypes.torch_dtype_from(dtype)
-
+        torch_dtype = _helpers.float_or_default(dtype, self.dtype)
         if axis is None:
             result = self._tensor.sum(dtype=torch_dtype)
         else:
             result = self._tensor.sum(dtype=torch_dtype, dim=axis)
 
         return result
 
+    @axis_out_keepdims_wrapper
+    def prod(self, axis=None, dtype=None, out=None, keepdims=NoValue,
+            initial=NoValue, where=NoValue):
+        if initial is not None or where is not None:
+            raise NotImplementedError
+
+        axis = _helpers.allow_only_single_axis(axis)
+
+        torch_dtype = _helpers.float_or_default(dtype, self.dtype)
+        if axis is None:
+            result = self._tensor.prod(dtype=torch_dtype)
+        else:
+            result = self._tensor.prod(dtype=torch_dtype, dim=axis)
+
+        return result
+
+
+    @axis_out_keepdims_wrapper
+    def std(self, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *,
+            where=NoValue):
+        if where is not None:
+            raise NotImplementedError
+
+        torch_dtype = _helpers.float_or_default(dtype, self.dtype)
+        tensor = self._tensor.to(torch_dtype)   # XXX: needed?
+
+        result = tensor.std(dim=axis, correction=ddof)
+
+        return result
+
+    @axis_out_keepdims_wrapper
+    def var(self, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *,
+            where=NoValue):
+        if where is not None:
+            raise NotImplementedError
+
+        torch_dtype = _helpers.float_or_default(dtype, self.dtype)
+        tensor = self._tensor.to(torch_dtype)   # XXX: needed?
+
+        result = tensor.var(dim=axis, correction=ddof)
+
+        return result
+
 
     ### indexing ###
     def __getitem__(self, *args, **kwds):

torch_np/_wrapper.py

@@ -307,20 +307,6 @@ def identity(n, dtype=None, *, like=None):
 ###### misc/unordered
 
 
-#YYY: pattern: initial=...
-@asarray_replacer()
-def prod(a, axis=None, dtype=None, out=None, keepdims=NoValue,
-         initial=NoValue, where=NoValue):
-    if initial is not None or where is not None:
-        raise NotImplementedError
-    if axis is None:
-        if keepdims is not None:
-            raise NotImplementedError
-        return torch.prod(a, dtype=dtype)
-    elif _util.is_sequence(axis):
-        raise NotImplementedError
-    return torch.prod(a, dim=axis, dtype=dtype, keepdim=bool(keepdims), out=out)
-
 
 
 @asarray_replacer()
@@ -639,13 +625,33 @@ def mean(a, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoValue)
     return arr.mean(axis=axis, dtype=dtype, out=out, keepdims=keepdims, where=where)
 
 
+#YYY: pattern: initial=...
+
 def sum(a, axis=None, dtype=None, out=None, keepdims=NoValue,
         initial=NoValue, where=NoValue):
     arr = asarray(a)
     return arr.sum(axis=axis, dtype=dtype, out=out, keepdims=keepdims,
                    initial=initial, where=where)
 
 
+def prod(a, axis=None, dtype=None, out=None, keepdims=NoValue,
+         initial=NoValue, where=NoValue):
+    arr = asarray(a)
+    return arr.prod(axis=axis, dtype=dtype, out=out, keepdims=keepdims,
+                   initial=initial, where=where)
+
+
+#YYY: pattern : ddof
+
+def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *, where=NoValue):
+    arr = asarray(a)
+    return arr.std(axis=axis, dtype=dtype, out=out, ddof=ddof, keepdims=keepdims, where=where)
+
+def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *, where=NoValue):
+    arr = asarray(a)
+    return arr.var(axis=axis, dtype=dtype, out=out, ddof=ddof, keepdims=keepdims, where=where)
+
+
 @asarray_replacer()
 def nanmean(a, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoValue):
     if where is not None:
@@ -663,30 +669,6 @@ def nanmean(a, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoVal
     return result
 
 
-# YYY: pattern : std, var
-@asarray_replacer()
-def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *, where=NoValue):
-    if where is not None:
-        raise NotImplementedError
-    if dtype is not None:
-        raise NotImplementedError 
-    if not torch.is_floating_point(a):
-        a = a * 1.0
-    return torch.std(a, axis, correction=ddof, keepdim=bool(keepdims), out=out)
-
-
-@asarray_replacer()
-def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *, where=NoValue):
-    if where is not None:
-        raise NotImplementedError
-    if dtype is not None:
-        raise NotImplementedError 
-    if not torch.is_floating_point(a):
-        a = a * 1.0
-    return torch.var(a, axis, correction=ddof, keepdim=bool(keepdims), out=out)
-
-
-
 @asarray_replacer()
 def argsort(a, axis=-1, kind=None, order=None):
     if order is not None:

torch_np/tests/test_reductions.py

@@ -593,3 +593,24 @@ def setup_method(self):
         self.allowed_axes =  [0, 1, 2, -1, -2,
                               (0, 1), (1, 0), (0, 1, 2), (1, -1, 0)]
 
+
+class TestProdGeneric(_GenericReductionsTestMixin, _GenericHasOutTestMixin):
+    def setup_method(self):
+        self.func = np.prod
+        self.allowed_axes =  [0, 1, 2, -1, -2,]
+#                              (0, 1), (1, 0), (0, 1, 2), (1, -1, 0)]
+
+
+class TestStdGeneric(_GenericReductionsTestMixin, _GenericHasOutTestMixin):
+    def setup_method(self):
+        self.func = np.std
+        self.allowed_axes =  [0, 1, 2, -1, -2,
+                              (0, 1), (1, 0), (0, 1, 2), (1, -1, 0)]
+
+
+class TestVarGeneric(_GenericReductionsTestMixin, _GenericHasOutTestMixin):
+    def setup_method(self):
+        self.func = np.var
+        self.allowed_axes =  [0, 1, 2, -1, -2,
+                              (0, 1), (1, 0), (0, 1, 2), (1, -1, 0)]
+