jax==0.8.2 breaks `is_jax_array` and subsequently `device`

[amacati]

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.

Relevant Issues

jax-ml/jax#26000
scipy/scipy#22246

[jakevdp]

I see - it looks like is_jax_array essentially uses issubclass(type(obj), jax.Array) in place of isinstance(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_compat to change to more standard instance checks, and use isinstance rather than issubclass. The fact that tracers were previously a subclass of Array was an implementation detail, and not something that downstream libraries should rely on.

[amacati]

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]

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]

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 ?

sounds like a good idea if we can sort the implementation out

[lucascolley]

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?

let us ping @crusaderky also

[crusaderky]

Speed is exactly the issue. Running the whole isinstance machinery, 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 in xp = array_namespace(x) at the top of scipy functions constituting the majority of the total runtime for small inputs (tens~hundreds of bytes).

See:

• ENH: cache helper functions #308
• ENH: speed up array_namespace #329
• ENH: Speed up array_namespace through caching scipy/scipy#23057

[ev-br]

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]

lots of people are away so a week across Christmas is way too short

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]

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.