Add benchmark for numba elemwise

Author: aseyboldt

tests/link/numba/test_elemwise.py

@@ -6,7 +6,7 @@
 import pytensor.tensor as at
 import pytensor.tensor.inplace as ati
 import pytensor.tensor.math as aem
-from pytensor import config
+from pytensor import config, function
 from pytensor.compile.ops import deep_copy_op
 from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.graph.basic import Constant
@@ -117,6 +117,25 @@ def test_Elemwise(inputs, input_vals, output_fn, exc):
         compare_numba_and_py(out_fg, input_vals)
 
 
+def test_elemwise_speed(benchmark):
+    x = at.dmatrix("y")
+    y = at.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",
     [