Make sum/prod/mean follow the same pattern other reductions do

torch_np/_helpers.py

@@ -131,3 +131,16 @@ 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, enforce_float=False):
+    """dtype helper for reductions."""
+    if dtype is None:
+        dtype = self_dtype
+    if dtype == _dtypes.dtype('bool'):
+        dtype = _dtypes.default_int_type()
+    if enforce_float:
+        if _dtypes.is_integer(dtype):
+            dtype = _dtypes.default_float_type()
+    torch_dtype = _dtypes.torch_dtype_from(dtype)
+    return torch_dtype

torch_np/_ndarray.py

@@ -86,6 +86,10 @@ def T(self):
     def real(self):
         return asarray(self._tensor.real)
 
+    @real.setter
+    def real(self, value):
+        self._tensor.real = asarray(value).get()
+
     @property
     def imag(self):
         try:
@@ -94,6 +98,10 @@ def imag(self):
             zeros = torch.zeros_like(self._tensor)
             return ndarray._from_tensor_and_base(zeros, None)
 
+    @imag.setter
+    def imag(self, value):
+        self._tensor.imag = asarray(value).get()
+
     # ctors
     def astype(self, dtype):
         newt = ndarray()
@@ -416,12 +424,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, enforce_float=True)
         if axis is None:
             result = self._tensor.mean(dtype=torch_dtype)
         else:
@@ -436,19 +439,53 @@ 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)
+kwargs = {"dim": axis} if axis is not None else {}
+return self._tensor.prod(dtype=torch_dtype, **kwargs)
+
+
+    @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, enforce_float=True)
+        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, enforce_float=True)
+        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:
@@ -775,6 +757,11 @@ def isscalar(a):
         return False
 
 
+def isclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
+    a = asarray(a).get()
+    b = asarray(a).get()
+    return asarray(torch.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan))
+
 ###### mapping from numpy API objects to wrappers from this module ######
 
 # All is in the mapping dict in _mapping.py

torch_np/testing/__init__.py

@@ -1,7 +1,7 @@
 from .utils import (assert_equal, assert_array_equal, assert_almost_equal,
-        assert_warns, assert_)
+        assert_warns, assert_, assert_allclose)
 from .utils import _gen_alignment_data
 
-from .testing import assert_allclose    # FIXME
+#from .testing import assert_allclose    # FIXME
 
 

torch_np/testing/utils.py

@@ -1169,7 +1169,7 @@ def _assert_valid_refcount(op):
 
 
 def assert_allclose(actual, desired, rtol=1e-7, atol=0, equal_nan=True,
-                    err_msg='', verbose=True):
+                    err_msg='', verbose=True, check_dtype=False):
     """
     Raises an AssertionError if two objects are not equal up to desired
     tolerance.
@@ -1226,14 +1226,17 @@ def assert_allclose(actual, desired, rtol=1e-7, atol=0, equal_nan=True,
 
     """
     __tracebackhide__ = True  # Hide traceback for py.test
-    import numpy as np
 
     def compare(x, y):
-        return np.core.numeric.isclose(x, y, rtol=rtol, atol=atol,
+        return np.isclose(x, y, rtol=rtol, atol=atol,
                                        equal_nan=equal_nan)
 
     actual, desired = asanyarray(actual), asanyarray(desired)
     header = f'Not equal to tolerance rtol={rtol:g}, atol={atol:g}'
+
+    if check_dtype:
+        assert actual.dtype == desired.dtype
+
     assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
                          verbose=verbose, header=header, equal_nan=equal_nan)