Group numba benchmark tests in same class

tests/link/numba/test_elemwise.py

@@ -130,25 +130,6 @@ def test_elemwise_runtime_broadcast():
     check_elemwise_runtime_broadcast(get_mode("NUMBA"))
 
 
-def test_elemwise_speed(benchmark):
-    x = pt.dmatrix("y")
-    y = pt.dvector("z")
-
-    out = np.exp(2 * x * y + y)
-
-    rng = np.random.default_rng(42)
-
-    x_val = rng.normal(size=(200, 500))
-    y_val = rng.normal(size=500)
-
-    func = function([x, y], out, mode="NUMBA")
-    func = func.vm.jit_fn
-    (out,) = func(x_val, y_val)
-    np.testing.assert_allclose(np.exp(2 * x_val * y_val + y_val), out)
-
-    benchmark(func, x_val, y_val)
-
-
 @pytest.mark.parametrize(
     "v, new_order",
     [
@@ -631,27 +612,55 @@ def test_Argmax(x, axes, exc):
         )
 
 
-@pytest.mark.parametrize("size", [(10, 10), (1000, 1000), (10000, 10000)])
-@pytest.mark.parametrize("axis", [0, 1])
-def test_logsumexp_benchmark(size, axis, benchmark):
-    X = pt.matrix("X")
-    X_max = pt.max(X, axis=axis, keepdims=True)
-    X_max = pt.switch(pt.isinf(X_max), 0, X_max)
-    X_lse = pt.log(pt.sum(pt.exp(X - X_max), axis=axis, keepdims=True)) + X_max
+def test_elemwise_out_type():
+    # Create a graph with an elemwise
+    # Ravel failes if the elemwise output type is reported incorrectly
+    x = pt.matrix()
+    y = (2 * x).ravel()
 
-    rng = np.random.default_rng(23920)
-    X_val = rng.normal(size=size)
+    # Pass in the input as mutable, to trigger the inplace rewrites
+    func = pytensor.function([pytensor.In(x, mutable=True)], y, mode="NUMBA")
 
-    X_lse_fn = pytensor.function([X], X_lse, mode="NUMBA")
+    # Apply it to a numpy array that is neither C or F contigous
+    x_val = np.broadcast_to(np.zeros((3,)), (6, 3))
+
+    assert func(x_val).shape == (18,)
 
-    # JIT compile first
-    res = X_lse_fn(X_val)
-    exp_res = scipy.special.logsumexp(X_val, axis=axis, keepdims=True)
-    np.testing.assert_array_almost_equal(res, exp_res)
-    benchmark(X_lse_fn, X_val)
 
+def test_scalar_loop():
+    a = float64("a")
+    scalar_loop = pytensor.scalar.ScalarLoop([a], [a + a])
 
-def test_fused_elemwise_benchmark(benchmark):
+    x = pt.tensor("x", shape=(3,))
+    elemwise_loop = Elemwise(scalar_loop)(3, x)
+
+    with pytest.warns(UserWarning, match="object mode"):
+        compare_numba_and_py(
+            ([x], [elemwise_loop]),
+            (np.array([1, 2, 3], dtype="float64"),),
+        )
+
+
+class TestsBenchmark:
+    def test_elemwise_speed(self, benchmark):
+        x = pt.dmatrix("y")
+        y = pt.dvector("z")
+
+        out = np.exp(2 * x * y + y)
+
+        rng = np.random.default_rng(42)
+
+        x_val = rng.normal(size=(200, 500))
+        y_val = rng.normal(size=500)
+
+        func = function([x, y], out, mode="NUMBA")
+        func = func.vm.jit_fn
+        (out,) = func(x_val, y_val)
+        np.testing.assert_allclose(np.exp(2 * x_val * y_val + y_val), out)
+
+        benchmark(func, x_val, y_val)
+
+    def test_fused_elemwise_benchmark(self, benchmark):
         rng = np.random.default_rng(123)
         size = 100_000
         x = pytensor.shared(rng.normal(size=size), name="x")
@@ -665,47 +674,36 @@ def test_fused_elemwise_benchmark(benchmark):
         func()
         benchmark(func)
 
+    @pytest.mark.parametrize("size", [(10, 10), (1000, 1000), (10000, 10000)])
+    @pytest.mark.parametrize("axis", [0, 1])
+    def test_logsumexp_benchmark(self, size, axis, benchmark):
+        X = pt.matrix("X")
+        X_max = pt.max(X, axis=axis, keepdims=True)
+        X_max = pt.switch(pt.isinf(X_max), 0, X_max)
+        X_lse = pt.log(pt.sum(pt.exp(X - X_max), axis=axis, keepdims=True)) + X_max
 
-def test_elemwise_out_type():
-    # Create a graph with an elemwise
-    # Ravel failes if the elemwise output type is reported incorrectly
-    x = pt.matrix()
-    y = (2 * x).ravel()
-
-    # Pass in the input as mutable, to trigger the inplace rewrites
-    func = pytensor.function([pytensor.In(x, mutable=True)], y, mode="NUMBA")
-
-    # Apply it to a numpy array that is neither C or F contigous
-    x_val = np.broadcast_to(np.zeros((3,)), (6, 3))
+        rng = np.random.default_rng(23920)
+        X_val = rng.normal(size=size)
 
-    assert func(x_val).shape == (18,)
+        X_lse_fn = pytensor.function([X], X_lse, mode="NUMBA")
 
+        # JIT compile first
+        res = X_lse_fn(X_val)
+        exp_res = scipy.special.logsumexp(X_val, axis=axis, keepdims=True)
+        np.testing.assert_array_almost_equal(res, exp_res)
+        benchmark(X_lse_fn, X_val)
 
-@pytest.mark.parametrize(
+    @pytest.mark.parametrize(
         "axis",
         (0, 1, 2, (0, 1), (0, 2), (1, 2), None),
         ids=lambda x: f"axis={x}",
-)
-@pytest.mark.parametrize(
+    )
+    @pytest.mark.parametrize(
         "c_contiguous",
         (True, False),
         ids=lambda x: f"c_contiguous={x}",
-)
-def test_numba_careduce_benchmark(axis, c_contiguous, benchmark):
+    )
+    def test_numba_careduce_benchmark(self, axis, c_contiguous, benchmark):
         return careduce_benchmark_tester(
             axis, c_contiguous, mode="NUMBA", benchmark=benchmark
         )
-
-
-def test_scalar_loop():
-    a = float64("a")
-    scalar_loop = pytensor.scalar.ScalarLoop([a], [a + a])
-
-    x = pt.tensor("x", shape=(3,))
-    elemwise_loop = Elemwise(scalar_loop)(3, x)
-
-    with pytest.warns(UserWarning, match="object mode"):
-        compare_numba_and_py(
-            ([x], [elemwise_loop]),
-            (np.array([1, 2, 3], dtype="float64"),),
-        )