Numba Pow: Fix failure with discrete integer exponents (#1758)

Workaround for: numba/numba#9554

Author: ricardoV94

pytensor/link/numba/dispatch/scalar.py

@@ -22,6 +22,7 @@
     Composite,
     Identity,
     Mul,
+    Pow,
     Reciprocal,
     ScalarOp,
     Second,
@@ -165,6 +166,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

@@ -189,6 +189,16 @@ def test_Softplus(dtype):
         )
 
 
+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")]
+    )
+
+
 def test_cython_obj_mode_fallback():
     """Test that unsupported cython signatures fallback to obj-mode"""