Repository: mars-project/mars

Issue:
[BUG] count_nonzero produces unexpected output for non-boolean tensor
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
``count_nonzero`` only produces the expected output for boolean arrays and integer arrays of 1s and 0s.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. *Python:* 3.7
2. *The version of Mars you use:* github master
3. *Versions of crucial packages, such as numpy, scipy and protobuf:* 1.18.1, 1.4.1, 3.11.2

**Expected behavior**
`count_nonzero` should count the number of non-zero values of any type which are not, 0, 0.0 or False etc.

```python
import numpy as np
import mars.tensor as mt
```

The following two work as expected:
```python
values = [False, True, False, True]
print(np.count_nonzero(np.array(values)))
print(mt.count_nonzero(mt.tensor(values)).execute())
# 2
# 2

values = [0, 1, 0, 1]
print(np.count_nonzero(np.array(values)))
print(mt.count_nonzero(mt.tensor(values)).execute())
# 2
# 2
```
but then we run into problems:
```python
values = [0, 2, 0, 3]
print(np.count_nonzero(np.array(values)))
print(mt.count_nonzero(mt.tensor(values)).execute())
# 2
# 5

values = [0.0, 0.2, 0.0, 0.3]
print(np.count_nonzero(np.array(values)))
print(mt.count_nonzero(mt.tensor(values)).execute())
# 2
# 0.5
```
It seems like `mars` is just doing a sum.




Discussion:
Thank you for your report, we will try to solve it as soon as possible.

Resolving patch:
diff --git a/mars/tensor/reduction/core.py b/mars/tensor/reduction/core.py
--- a/mars/tensor/reduction/core.py
+++ b/mars/tensor/reduction/core.py
@@ -271,6 +271,21 @@ def execute_agg(cls, ctx, op):
 
             ctx[out.key] = ret.astype(ret.dtype, order=out.order.value, copy=False)
 
+    @classmethod
+    def execute_one_chunk(cls, ctx, op):
+        cls.execute_agg(ctx, op)
+
+    @classmethod
+    def execute(cls, ctx, op):
+        if op.stage == OperandStage.map:
+            return cls.execute_map(ctx, op)
+        elif op.stage == OperandStage.combine:
+            return cls.execute_combine(ctx, op)
+        elif op.stage == OperandStage.agg:
+            return cls.execute_agg(ctx, op)
+        else:
+            return cls.execute_one_chunk(ctx, op)
+
 
 class TensorArgReductionMixin(TensorReductionMixin):
     __slots__ = ()
@@ -522,15 +537,6 @@ def _rewrite_stage(self, stage):
             return OperandStage.agg
         return stage
 
-    @classmethod
-    def execute(cls, ctx, op):
-        if op.stage == OperandStage.map:
-            return cls.execute_map(ctx, op)
-        elif op.stage == OperandStage.combine:
-            return cls.execute_combine(ctx, op)
-        else:
-            return cls.execute_agg(ctx, op)
-
 
 class TensorCumReduction(TensorHasInput):
     _input = KeyField('input')
diff --git a/mars/tensor/reduction/count_nonzero.py b/mars/tensor/reduction/count_nonzero.py
--- a/mars/tensor/reduction/count_nonzero.py
+++ b/mars/tensor/reduction/count_nonzero.py
@@ -51,6 +51,14 @@ def execute_map(cls, ctx, op):
     def execute_agg(cls, ctx, op):
         return TensorSum.execute_agg(ctx, op)
 
+    @classmethod
+    def execute_one_chunk(cls, ctx, op):
+        a = ctx[op.inputs[0].key]
+        (inp,), device_id, xp = as_same_device(
+            [a], device=op.device, ret_extra=True)
+        with device(device_id):
+            ctx[op.outputs[0].key] = xp.count_nonzero(inp, axis=op.axis)
+
 
 def count_nonzero(a, axis=None, combine_size=None):
     """


Test patch:
diff --git a/mars/tensor/reduction/tests/test_reduction_execute.py b/mars/tensor/reduction/tests/test_reduction_execute.py
--- a/mars/tensor/reduction/tests/test_reduction_execute.py
+++ b/mars/tensor/reduction/tests/test_reduction_execute.py
@@ -173,6 +173,12 @@ def testVarExecution(self):
         raw1 = np.random.random((20, 25))
         raw2 = np.random.randint(10, size=(20, 25))
 
+        arr0 = tensor(raw1, chunk_size=25)
+
+        res1 = self.executor.execute_tensor(arr0.var())
+        expected1 = raw1.var()
+        self.assertTrue(np.allclose(res1[0], expected1))
+
         arr1 = tensor(raw1, chunk_size=3)
 
         res1 = self.executor.execute_tensor(arr1.var())
@@ -465,6 +471,13 @@ def testOutCumReductionExecution(self):
     def testCountNonzeroExecution(self):
         raw = [[0, 1, 7, 0, 0], [3, 0, 0, 2, 19]]
 
+        arr = tensor(raw, chunk_size=5)
+        t = count_nonzero(arr)
+
+        res = self.executor.execute_tensor(t)[0]
+        expected = np.count_nonzero(raw)
+        np.testing.assert_equal(res, expected)
+
         arr = tensor(raw, chunk_size=2)
         t = count_nonzero(arr)