FIX (0.21) make count_nonzero dtype invariant wrt axis (#12341)

jnothman · qinhanmin2014 · commit 98357ec4d4e4 · 2018-10-14T21:12:02.000+08:00
diff --git a/sklearn/utils/sparsefuncs.py b/sklearn/utils/sparsefuncs.py
@@ -467,7 +467,8 @@ def count_nonzero(X, axis=None, sample_weight=None):
     elif axis == 1:
         out = np.diff(X.indptr)
         if sample_weight is None:
-            return out
+            # astype here is for consistency with axis=0 dtype
+            return out.astype('intp')
         return out * sample_weight
     elif axis == 0:
         if sample_weight is None:
diff --git a/sklearn/utils/tests/test_sparsefuncs.py b/sklearn/utils/tests/test_sparsefuncs.py
@@ -443,6 +443,19 @@ def test_count_nonzero():
     assert_raises(TypeError, count_nonzero, X_csc)
     assert_raises(ValueError, count_nonzero, X_csr, axis=2)
 
+    assert (count_nonzero(X_csr, axis=0).dtype ==
+            count_nonzero(X_csr, axis=1).dtype)
+    assert (count_nonzero(X_csr, axis=0, sample_weight=sample_weight).dtype ==
+            count_nonzero(X_csr, axis=1, sample_weight=sample_weight).dtype)
+
+    # Check dtypes with large sparse matrices too
+    X_csr.indices = X_csr.indices.astype(np.int64)
+    X_csr.indptr = X_csr.indptr.astype(np.int64)
+    assert (count_nonzero(X_csr, axis=0).dtype ==
+            count_nonzero(X_csr, axis=1).dtype)
+    assert (count_nonzero(X_csr, axis=0, sample_weight=sample_weight).dtype ==
+            count_nonzero(X_csr, axis=1, sample_weight=sample_weight).dtype)
+
 
 def test_csc_row_median():
     # Test csc_row_median actually calculates the median.
