Repository navigation
jax==0.8.2 breaks is_jax_array and subsequently device #368
Description
Jax just released its version 0.8.2 https://github.com/jax-ml/jax/releases/tag/jax-v0.8.2, and it causes is_jax_array to no longer work for jit compiled arrays. It's probably caused be this change:
jax's Tracer no longer inherits from jax.Array at runtime. However, jax.Array now uses a custom metaclass such isinstance(x, Array) is true if an object x represents a traced Array. Only some Tracers represent Arrays, so it is not correct for Tracer to inherit from Array.
Example
import array_api_compat as xpc
import jax
import jax.numpy as jnp
@jax.jit
def fn(x):
print(xpc.is_jax_array(x)) # False
return jnp.zeros(x.shape, device=xpc.device(x))
x = jnp.array([1.0, 2.0, 3.0])
print(xpc.is_jax_array(x)) # True
y = fn(x)which yields
Traceback (most recent call last):
File "/home/mschuck/repos/proto/.venv/lib/python3.14/site-packages/jax/_src/core.py", line 1071, in __getattr__
attr = getattr(self.aval, name)
AttributeError: 'ShapedArray' object has no attribute 'device'. Did you mean: 'to_device'?
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/mschuck/repos/proto/proto/rotation_grad.py", line 28, in <module>
y = fn(x)
File "/home/mschuck/repos/proto/proto/rotation_grad.py", line 24, in fn
return jnp.zeros(x.shape, device=xpc.device(x))
~~~~~~~~~~^^^
File "/home/mschuck/repos/proto/.venv/lib/python3.14/site-packages/array_api_compat/common/_helpers.py", line 764, in device
return x.device # pyright: ignore
^^^^^^^^
AttributeError: DynamicJaxprTracer has no attribute device. Did you mean: 'devices'?
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
Activity
jakevdp commentedon Dec 19, 2025
I see - it looks like
is_jax_arrayessentially usesissubclass(type(obj), jax.Array)in place ofisinstance(obj, jax.Array). The JAX team ensured that the latter works with the inheritence change, but the former does not.My suggestion would be for
array_api_compatto change to more standard instance checks, and useisinstancerather thanissubclass. The fact that tracers were previously a subclass of Array was an implementation detail, and not something that downstream libraries should rely on.amacati commentedon Dec 19, 2025
Yes, this hopefully is an easy fix. Given that scipy is releasing substantially more array API capabilities with 1.17, we should get a patch out before the final release if possible and update the version we vendor with the release. Is that possible @lucascolley ?
amacati commentedon Dec 19, 2025
Looking at it, it seems like the design is using hashable objects instead of the array itself to enable caching the result. Since Arrays are not hashable, their type name is used instead. Is there a way to link jax Arrays to tracers by their types only, or is that impossible with the latest changes in JAX?
lucascolley commentedon Dec 19, 2025
sounds like a good idea if we can sort the implementation out
lucascolley commentedon Dec 19, 2025
let us ping @crusaderky also
crusaderky commentedon Dec 19, 2025
Speed is exactly the issue. Running the whole
isinstancemachinery, plus special hacks for zero gradient arrays (jax-ml/jax#20620), on each array instance was up to 5x more expensive than getting a cache hit on types. That culminated inxp = array_namespace(x)at the top of scipy functions constituting the majority of the total runtime for small inputs (tens~hundreds of bytes).See:
array_namespace#329array_namespacethrough caching scipy/scipy#23057ev-br commentedon Dec 27, 2025
Please leave some time for folks to take a look. The Christmas/New Year season is expected to be slow, lots of people are away so a week across Christmas is way too short.
lucascolley commentedon Dec 27, 2025
point taken, however there is the consideration of a potential backport for SciPy 1.17.0 which pushes in the opposite direction.
I judged that my review and approval on top of the input from Jake was sufficient over in gh-369, but we don't have any maintainer documentation to tell either way. Maybe something for the agenda of the next community meeting if we feel like we should have a different process.
ev-br commentedon Dec 28, 2025
Not sure if the existence of a Christmas break is something to discuss in the community meeting. Also I don't think we need more process TBH. I think we're in the clear here as is.
Anyhow, please feel free to push the button on the release PR if you think it should be pressed on. I typically run array-api tests with major backends locally + test scipy and scikit-learn before pressing the button.