Merge pull request numpy#9638 from juliantaylor/nonzero-dtype

BUG: ensure consistent result dtype of count_nonzero

Author: ahaldane

numpy/core/numeric.py

@@ -444,7 +444,7 @@ def count_nonzero(a, axis=None):
     counts = np.apply_along_axis(multiarray.count_nonzero, axis[0], a)
 
     if axis.size == 1:
-        return counts
+        return counts.astype(np.intp, copy=False)
     else:
         # for subsequent axis numbers, that number decreases
         # by one in this new 'counts' array if it was larger

numpy/core/tests/test_numeric.py

@@ -1026,6 +1026,10 @@ def test_count_nonzero_axis_all_dtypes(self):
         # either integer or tuple arguments for axis
         msg = "Mismatch for dtype: %s"
 
+        def assert_equal_w_dt(a, b, err_msg):
+            assert_equal(a.dtype, b.dtype, err_msg=err_msg)
+            assert_equal(a, b, err_msg=err_msg)
+
         for dt in np.typecodes['All']:
             err_msg = msg % (np.dtype(dt).name,)
 
@@ -1045,13 +1049,13 @@ def test_count_nonzero_axis_all_dtypes(self):
                     m[1, 0] = '1970-01-12'
                     m = m.astype(dt)
 
-                expected = np.array([2, 0, 0])
-                assert_equal(np.count_nonzero(m, axis=0),
-                             expected, err_msg=err_msg)
+                expected = np.array([2, 0, 0], dtype=np.intp)
+                assert_equal_w_dt(np.count_nonzero(m, axis=0),
+                                  expected, err_msg=err_msg)
 
-                expected = np.array([1, 1, 0])
-                assert_equal(np.count_nonzero(m, axis=1),
-                             expected, err_msg=err_msg)
+                expected = np.array([1, 1, 0], dtype=np.intp)
+                assert_equal_w_dt(np.count_nonzero(m, axis=1),
+                                  expected, err_msg=err_msg)
 
                 expected = np.array(2)
                 assert_equal(np.count_nonzero(m, axis=(0, 1)),
@@ -1066,13 +1070,13 @@ def test_count_nonzero_axis_all_dtypes(self):
                 # setup is slightly different for this dtype
                 m = np.array([np.void(1)] * 6).reshape((2, 3))
 
-                expected = np.array([0, 0, 0])
-                assert_equal(np.count_nonzero(m, axis=0),
-                             expected, err_msg=err_msg)
+                expected = np.array([0, 0, 0], dtype=np.intp)
+                assert_equal_w_dt(np.count_nonzero(m, axis=0),
+                                  expected, err_msg=err_msg)
 
-                expected = np.array([0, 0])
-                assert_equal(np.count_nonzero(m, axis=1),
-                             expected, err_msg=err_msg)
+                expected = np.array([0, 0], dtype=np.intp)
+                assert_equal_w_dt(np.count_nonzero(m, axis=1),
+                                  expected, err_msg=err_msg)
 
                 expected = np.array(0)
                 assert_equal(np.count_nonzero(m, axis=(0, 1)),