Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

jax.numpy.argmin/jax.numpy.argmax fail with nans

See original GitHub issue

Noticed odd behavior (colab repro) and was about to fix/submit a PR when I saw this comment:

# TODO(mattjj): redo this lowering with a call to variadic lax.reduce

Wasn’t sure what @mattjj had in mind, but happy to have a look at that as well.

Issue Analytics

State:
Created 3 years ago
Comments:13 (9 by maintainers)

Top GitHub Comments

1reaction

jakevdpcommented, Jun 1, 2021

Thanks - looks like it’s only a problem on GPU, not on CPU or TPU. This suggests it may be an XLA issue.

0reactions

hawkinspcommented, Jun 1, 2021

I think this was already fixed by https://github.com/google/jax/pull/6764 . That PR isn’t in a jax release yet, though, so to see the fix you’d need to use jax from Github head.

(This both was and was not an XLA issue. It’s an XLA issue that GPU acts differently, because XLA/GPU lacks support for variadic reductions, forcing us to use a different code path on GPU. The fallback path handling NaNs differently was a JAX issue.)

Top Results From Across the Web

numpy.nanargmax() - JAX documentation - Read the Docs

argmax returns -1 for all-NaN slices and does not raise an error. Original docstring below. NaNs. For all-NaN slices ValueError is raised. Warning:...