Workaround numba/numba#9554

Author: ricardoV94

pytensor/link/numba/dispatch/scalar.py

@@ -23,6 +23,7 @@
     Composite,
     Identity,
     Mul,
+    Pow,
     Reciprocal,
     ScalarOp,
     Second,
@@ -171,6 +172,23 @@ def {binary_op_name}({input_signature}):
     return nary_fn
 
 
+@register_funcify_and_cache_key(Pow)
+def numba_funcify_Pow(op, node, **kwargs):
+    pow_dtype = node.inputs[1].type.dtype
+    if pow_dtype.startswith("int"):
+        # Numba power fails when exponents are non 64-bit discrete integers and fasthmath=True
+        # https://github.com/numba/numba/issues/9554
+
+        def pow(x, y):
+            return x ** np.asarray(y, dtype=np.int64).item()
+    else:
+
+        def pow(x, y):
+            return x**y
+
+    return numba_basic.numba_njit(pow), scalar_op_cache_key(op)
+
+
 @register_funcify_and_cache_key(Add)
 def numba_funcify_Add(op, node, **kwargs):
     nary_add_fn = binary_to_nary_func(node.inputs, "add", "+")

tests/link/numba/test_scalar.py

@@ -184,3 +184,13 @@ def test_Softplus(dtype):
             strict=True,
             err_msg=f"Failed for value {value}",
         )
+
+
+def test_discrete_power():
+    # Test we don't fail to compile power with discrete exponents due to https://github.com/numba/numba/issues/9554
+    x = pt.scalar("x", dtype="float64")
+    exponent = pt.scalar("exponent", dtype="int8")
+    out = pt.power(x, exponent)
+    compare_numba_and_py(
+        [x, exponent], [out], [np.array(0.5), np.array(2, dtype="int8")]
+    )