`pt.isnan` does not work in NUMBA mode

[jessegrabowski]

Description

See below:

import pytensor
import pytensor.tensor as pt

y = pt.dvector('y')
nan_mask = pt.isnan(y)
funcs = [pytensor.function([y], nan_mask, mode=mode) for mode in [None, 'JAX', 'NUMBA']]

[f(np.array([np.nan, 1, 2, 3])) for f in funcs]
# C:            [array([ True, False, False, False]),
# JAX:         Array([ True, False, False, False], dtype=bool),
# NUMBA:  array([False, False, False, False])]

[ricardoV94]

It's the behavior of fastmath:

import numpy as np
import numba

@numba.njit(fastmath=True)
def isnan(x):
    return np.isnan(x)

isnan(np.array([np.nan, 1, 2, 3]))  # array([False, False, False, False])

numba/numba#9383
numba/numba#2919

[ricardoV94]

I guess we want to manually disable fastmath for the scalar Op IsNan in our case