ENH: added axis param for np.count_nonzero

Closes gh-391.

2 files changed, 24 insertions, 3 deletions

diff --git a/benchmarks/benchmarks/bench_core.py b/benchmarks/benchmarks/bench_core.py
index b595ae469..6701917cc 100644
--- a/benchmarks/benchmarks/bench_core.py
+++ b/benchmarks/benchmarks/bench_core.py
@@ -106,3 +106,27 @@ class CorrConv(Benchmark):
 
     def time_convolve(self, size1, size2, mode):
         np.convolve(self.x1, self.x2, mode=mode)
+
+
+class CountNonzero(Benchmark):
+    param_names = ['numaxes', 'size', 'dtype']
+    params = [
+        [1, 2, 3],
+        [100, 10000, 1000000],
+        [bool, int, str, object]
+    ]
+
+    def setup(self, numaxes, size, dtype):
+        self.x = np.empty(shape=(
+            numaxes, size), dtype=dtype)
+
+    def time_count_nonzero(self, numaxes, size, dtype):
+        np.count_nonzero(self.x)
+
+    def time_count_nonzero_axis(self, numaxes, size, dtype):
+        np.count_nonzero(self.x, axis=self.x.ndim - 1)
+
+    def time_count_nonzero_multi_axis(self, numaxes, size, dtype):
+        if self.x.ndim >= 2:
+            np.count_nonzero(self.x, axis=(
+                self.x.ndim - 1, self.x.ndim - 2))
diff --git a/benchmarks/benchmarks/bench_ufunc.py b/benchmarks/benchmarks/bench_ufunc.py
index 8f821ce08..0140718e0 100644
--- a/benchmarks/benchmarks/bench_ufunc.py
+++ b/benchmarks/benchmarks/bench_ufunc.py
@@ -67,9 +67,6 @@ class Custom(Benchmark):
     def time_nonzero(self):
         np.nonzero(self.b)
 
-    def time_count_nonzero(self):
-        np.count_nonzero(self.b)
-
     def time_not_bool(self):
         (~self.b)