BUG: sum() of a bool array is integer

Author: ev-br

torch_np/_helpers.py

@@ -133,11 +133,14 @@ def to_tensors(*inputs):
             for value in inputs])
 
 
-def float_or_default(dtype, self_dtype):
+def float_or_default(dtype, self_dtype, enforce_float=False):
     """dtype helper for reductions."""
     if dtype is None:
         dtype = self_dtype
-    if _dtypes.is_integer(dtype):
-        dtype = _dtypes.default_float_type()
+    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

@@ -416,7 +416,7 @@ def mean(self, axis=None, dtype=None, out=None, keepdims=NoValue, *, where=NoVal
         if where is not None:
             raise NotImplementedError
 
-        torch_dtype = _helpers.float_or_default(dtype, self.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:
@@ -462,7 +462,7 @@ def std(self, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *,
         if where is not None:
             raise NotImplementedError
 
-        torch_dtype = _helpers.float_or_default(dtype, self.dtype)
+        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)
@@ -475,7 +475,7 @@ def var(self, axis=None, dtype=None, out=None, ddof=0, keepdims=NoValue, *,
         if where is not None:
             raise NotImplementedError
 
-        torch_dtype = _helpers.float_or_default(dtype, self.dtype)
+        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)

torch_np/tests/test_reductions.py

@@ -328,6 +328,16 @@ def test_sum_stability(self):
         assert_allclose((a / 10.).sum() - a.size / 10., 0., atol=1.5e-13,
                         check_dtype=False)
 
+    def test_sum_boolean(self):
+        a = (np.arange(7) % 2 == 0)
+        res = a.sum()
+        assert_equal(res, 4)
+
+        res_float = a.sum(dtype=np.float64)
+        assert_allclose(res_float, 4.0, atol=1e-15)
+        assert res_float.dtype == 'float64'
+
+
     @pytest.mark.xfail(reason="dtype(value) needs implementing")
     def test_sum_dtypes(self):
         for dt in (int, np.float16, np.float32, np.float64):