Override NUMBA fastmath flag in IsNan dispatch

Author: ricardoV94

pytensor/link/numba/dispatch/scalar.py

@@ -23,6 +23,7 @@
     Clip,
     Composite,
     Identity,
+    IsNan,
     Mul,
     Reciprocal,
     ScalarOp,
@@ -137,7 +138,8 @@ def {scalar_op_fn_name}({', '.join(input_names)}):
 
     return numba_basic.numba_njit(
         signature,
-        fastmath=config.numba__fastmath,
+        # numba always returns False if fastmath=True  # https://github.com/numba/numba/issues/9383
+        fastmath=False if isinstance(op, IsNan) else config.numba__fastmath,
         # Functions that call a function pointer can't be cached
         cache=False,
     )(scalar_op_fn)

tests/link/numba/test_scalar.py

@@ -9,6 +9,7 @@
 from pytensor.graph.basic import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.scalar.basic import Composite
+from pytensor.tensor import tensor
 from pytensor.tensor.elemwise import Elemwise
 from tests.link.numba.test_basic import compare_numba_and_py, set_test_value
 
@@ -140,3 +141,14 @@ def test_reciprocal(v, dtype):
             if not isinstance(i, SharedVariable | Constant)
         ],
     )
+
+
+@pytest.mark.parametrize("dtype", ("complex64", "float64", "float32"))
+def test_isnan(dtype):
+    # Testing with tensor just to make sure Elemwise does not revert the scalar behavior of fastmath
+    x = tensor(shape=(2,), dtype=dtype)
+    out = pt.isnan(x)
+    compare_numba_and_py(
+        ([x], [out]),
+        [np.array([1, 0], dtype=dtype)],
+    )