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