Work-around numba pow failure was not robust enough

Author: ricardoV94

pytensor/link/numba/dispatch/scalar.py

@@ -170,18 +170,17 @@ def {binary_op_name}({input_signature}):
 @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
 
-        def pow(x, y):
-            return x**y
+    # Numba power fails when exponents are discrete integers and fasthmath=True
+    # https://github.com/numba/numba/issues/9554
+    fastmath = False if pow_dtype.startswith("int") else None
 
-    return numba_basic.numba_njit(pow), scalar_op_cache_key(op)
+    return numba_basic.numba_njit(pow, fastmath=fastmath), scalar_op_cache_key(
+        op, cache_version=1
+    )
 
 
 @register_funcify_and_cache_key(Add)

tests/link/numba/test_scalar.py

@@ -185,11 +185,14 @@ 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")
+    x1 = pt.scalar("x1", dtype="float64")
+    x2 = pt.scalar("x2", dtype="int16")
     exponent = pt.scalar("exponent", dtype="int8")
-    out = pt.power(x, exponent)
+    outs = [pt.power(x1, exponent), pt.power(x2, exponent)]
     compare_numba_and_py(
-        [x, exponent], [out], [np.array(0.5), np.array(2, dtype="int8")]
+        [x1, x2, exponent],
+        outs,
+        [np.array(0.5), np.array(2, dtype="int16"), np.array(2, dtype="int8")],
     )