Fix is_jax_array for jax>=0.8.2

Author: amacati

array_api_compat/common/_helpers.py

@@ -235,7 +235,11 @@ def is_jax_array(x: object) -> TypeIs[jax.Array]:
     is_pydata_sparse_array
     """
     cls = cast(Hashable, type(x))
-    return _issubclass_fast(cls, "jax", "Array") or _is_jax_zero_gradient_array(x)
+    return (
+        _issubclass_fast(cls, "jax", "Array")
+        or _issubclass_fast(cls, "jax.core", "Tracer")
+        or _is_jax_zero_gradient_array(x)
+    )
 
 
 def is_pydata_sparse_array(x: object) -> TypeIs[sparse.SparseArray]:

tests/test_common.py

@@ -56,6 +56,15 @@ def test_is_xp_array(library, func):
     assert is_array_api_obj(x)
 
 
+def test_is_jax_array_jitted():
+    import jax
+    import jax.numpy as jnp
+
+    x = jnp.asarray([1, 2, 3])
+    assert is_jax_array(x)
+    assert jax.jit(lambda y: is_jax_array(y))(x)
+
+
 @pytest.mark.parametrize('library', is_namespace_functions.keys())
 @pytest.mark.parametrize('func', is_namespace_functions.values())
 def test_is_xp_namespace(library, func):