Skip to content

Commit

Permalink
Override NUMBA fastmath flag in IsNan dispatch
Browse files Browse the repository at this point in the history
  • Loading branch information
ricardoV94 committed Jan 2, 2025
1 parent 8cc489b commit c59ba2e
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 1 deletion.
4 changes: 3 additions & 1 deletion pytensor/link/numba/dispatch/scalar.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
Clip,
Composite,
Identity,
IsNan,
Mul,
Reciprocal,
ScalarOp,
Expand Down Expand Up @@ -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)
Expand Down
12 changes: 12 additions & 0 deletions tests/link/numba/test_scalar.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)],
)

0 comments on commit c59ba2e

Please sign in to comment.