Merge pull request numpy#27517 from eendebakpt/nonzero_benchmarks

BENCH: Add benchmarks for np.non_zero

benchmarks/benchmarks/bench_core.py

@@ -170,9 +170,34 @@ def time_count_nonzero_multi_axis(self, numaxes, size, dtype):
                 self.x.ndim - 1, self.x.ndim - 2))
 
 
+class Nonzero(Benchmark):
+    params = [
+        [bool, np.uint8, np.uint64, np.int64, np.float32, np.float64],
+        [(1_000_000,), (1000, 1000), (100, ), (2, )]
+    ]
+    param_names = ["dtype", "shape"]
+
+    def setup(self, dtype, size):
+        self.x = np.random.randint(0, 3, size=size).astype(dtype)
+        self.x_sparse = np.zeros(size).astype(dtype)
+        self.x_sparse[1] = 1
+        self.x_sparse[-1] = 1
+        self.x_dense = np.ones(size).astype(dtype)
+
+    def time_nonzero(self, dtype, size):
+        np.nonzero(self.x)
+
+    def time_nonzero_sparse(self, dtype, size):
+        np.nonzero(self.x_sparse)
+
+    def time_nonzero_dense(self, dtype, size):
+        np.nonzero(self.x_dense)
+
+
 class PackBits(Benchmark):
     param_names = ['dtype']
     params = [[bool, np.uintp]]
+
     def setup(self, dtype):
         self.d = np.ones(10000, dtype=dtype)
         self.d2 = np.ones((200, 1000), dtype=dtype)