From 8a730beff1d3e74b5eac0fed1c8bf1b21e5b175c Mon Sep 17 00:00:00 2001
From: Ralf Gommers <ralf.gommers@gmail.com>
Date: Sun, 27 Aug 2023 21:47:32 +0200
Subject: [PATCH] BENCH: fix issues with `operator` benchmarks in
 `bench_ufunc.py`

---
 benchmarks/benchmarks/bench_ufunc.py | 50 +++++++++++++++++++++-------
 1 file changed, 38 insertions(+), 12 deletions(-)

diff --git a/benchmarks/benchmarks/bench_ufunc.py b/benchmarks/benchmarks/bench_ufunc.py
index 8b18663ff926..45f2f312f4bb 100644
--- a/benchmarks/benchmarks/bench_ufunc.py
+++ b/benchmarks/benchmarks/bench_ufunc.py
@@ -159,23 +159,49 @@ def time_ndarray__0d__(self, methname, npdtypes):
 class MethodsV1(Benchmark):
     """ Benchmark for the methods which take an argument
     """
-    params = [['__and__', '__add__', '__eq__', '__floordiv__', '__ge__',
-               '__gt__', '__le__', '__lt__', '__matmul__',
-               '__mod__', '__mul__', '__ne__', '__or__',
-               '__pow__', '__sub__', '__truediv__', '__xor__'],
+    params = [['__add__', '__eq__', '__ge__', '__gt__', '__le__',
+               '__lt__', '__matmul__', '__mul__', '__ne__',
+               '__pow__', '__sub__', '__truediv__'],
               TYPES1]
     param_names = ['methods', 'npdtypes']
     timeout = 10
 
     def setup(self, methname, npdtypes):
-        if (
-            npdtypes.startswith("complex")
-                and methname in ["__floordiv__", "__mod__"]
-        ) or (
-            not npdtypes.startswith("int")
-            and methname in ["__and__", "__or__", "__xor__"]
-        ):
-            raise NotImplementedError  # skip
+        values = get_squares_().get(npdtypes)
+        self.xargs = [values[0], values[1]]
+        if np.issubdtype(npdtypes, np.inexact):
+            # avoid overflow in __pow__/__matmul__ for low-precision dtypes
+            self.xargs[1] *= 0.01
+
+    def time_ndarray_meth(self, methname, npdtypes):
+        getattr(operator, methname)(*self.xargs)
+
+
+class MethodsV1IntOnly(Benchmark):
+    """ Benchmark for the methods which take an argument
+    """
+    params = [['__and__', '__or__', '__xor__'],
+              ['int16', 'int32', 'int64']]
+    param_names = ['methods', 'npdtypes']
+    timeout = 10
+
+    def setup(self, methname, npdtypes):
+        values = get_squares_().get(npdtypes)
+        self.xargs = [values[0], values[1]]
+
+    def time_ndarray_meth(self, methname, npdtypes):
+        getattr(operator, methname)(*self.xargs)
+
+
+class MethodsV1NoComplex(Benchmark):
+    """ Benchmark for the methods which take an argument
+    """
+    params = [['__floordiv__', '__mod__'],
+              [dt for dt in TYPES1 if not dt.startswith('complex')]]
+    param_names = ['methods', 'npdtypes']
+    timeout = 10
+
+    def setup(self, methname, npdtypes):
         values = get_squares_().get(npdtypes)
         self.xargs = [values[0], values[1]]