Incorrect handling of nan in in_top_k of Jax backend (regression from keras 3.3.3)

[michaelpradel]

The top_in_k function in the Jax backend provides an unexpected result.

Example to reproduce the problem:

import jax.numpy as jnp
from numpy import nan
from keras.src.backend.jax.math import in_top_k

r = in_top_k(targets=jnp.array([1, 0]), predictions=jnp.array([[.1, nan, .5], [.3, .2, .5]]), k=2)
print(r)

With keras 3.3.3, I'm getting the expected outcome:
[False True]

However, keras 3.4.1 gives this:
[ True True]

The new behavior is unexpected because nan shouldn't be be considered as large probability in the prediction.

As a first step to debug it: The change in behavior has been introduced by #19814.

[sachinprasadhs]

I was able to replicate the reported behavior with JAX backend, attaching the Gist for reference
.

Tensorflow backend is working as expected in the latest Keras version.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[michaelpradel]

Thanks for fixing this so quickly!