Skip to content

ENH: _lib: deobfuscate jax.jit crash in _asarray #22219

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 31, 2024

Conversation

crusaderky
Copy link
Contributor

@crusaderky crusaderky commented Dec 31, 2024

>>> import jax
>>> import jax.numpy as xp
>>> from scipy._lib._array_api import _asarray
>>> a = xp.asarray([0])
>>> def f(a):
...     return _asarray(a, check_finite=True)
>>> jax.jit(f)(a)

Before:

ValueError: array must not contain infs or NaNs

After:

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function kmeans at /home/crusaderky/github/scipy/build-install/lib/python3.12/site-packages/scipy/cluster/vq.py:332 for jit. This concrete value was not available in Python because it depends on the value of the argument obs.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

Outside of JAX, it makes perfect sense to me that if you call _asarray(..., check_finite=True) on an array of strings you ought to get a meaningful TypeError.

@github-actions github-actions bot added scipy._lib array types Items related to array API support and input array validation (see gh-18286) enhancement A new feature or improvement labels Dec 31, 2024
Copy link
Member

@lucascolley lucascolley left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

makes sense, thanks!

@lucascolley lucascolley added this to the 1.16.0 milestone Dec 31, 2024
@lucascolley lucascolley merged commit 4758525 into scipy:main Dec 31, 2024
37 of 38 checks passed
@crusaderky crusaderky deleted the jax_asarray branch December 31, 2024 16:12
@crusaderky
Copy link
Contributor Author

Related: data-apis/array-api-compat#225

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
array types Items related to array API support and input array validation (see gh-18286) enhancement A new feature or improvement scipy._lib
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants