Fix and test array_api_obj, is_writable_array, is_lazy_array

Author: amacati

array_api_compat/common/_helpers.py

@@ -300,6 +300,7 @@ def _is_array_api_cls(cls: type) -> bool:
         or _issubclass_fast(cls, "sparse", "SparseArray")
         # TODO: drop support for jax<0.4.32 which didn't have __array_namespace__
         or _issubclass_fast(cls, "jax", "Array")
+        or _issubclass_fast(cls, "jax.core", "Tracer")
     )
 
 
@@ -938,6 +939,7 @@ def _is_writeable_cls(cls: type) -> bool | None:
     if (
         _issubclass_fast(cls, "numpy", "generic")
         or _issubclass_fast(cls, "jax", "Array")
+        or _issubclass_fast(cls, "jax.core", "Tracer")
         or _issubclass_fast(cls, "sparse", "SparseArray")
     ):
         return False
@@ -977,6 +979,7 @@ def _is_lazy_cls(cls: type) -> bool | None:
         return False
     if (
         _issubclass_fast(cls, "jax", "Array")
+        or _issubclass_fast(cls, "jax.core", "Tracer")
         or _issubclass_fast(cls, "dask.array", "Array")
         or _issubclass_fast(cls, "ndonnx", "Array")
     ):

tests/test_common.py

@@ -56,15 +56,6 @@ def test_is_xp_array(library, func):
     assert is_array_api_obj(x)
 
 
-def test_is_jax_array_jitted():
-    jax = pytest.importorskip("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):

tests/test_jax.py

@@ -1,7 +1,14 @@
 from numpy.testing import assert_equal
 import pytest
 
-from array_api_compat import device, to_device
+from array_api_compat import (
+    device,
+    to_device,
+    is_jax_array,
+    is_lazy_array,
+    is_array_api_obj,
+    is_writeable_array,
+)
 
 try:
     import jax
@@ -13,7 +20,7 @@
 
 
 @pytest.mark.parametrize(
-    "func", 
+    "func",
     [
         lambda x: jnp.zeros(1, device=device(x)),
         lambda x: jnp.zeros_like(jnp.ones(1, device=device(x))),
@@ -26,7 +33,7 @@
             ),
         ),
         lambda x: to_device(jnp.zeros(1), device(x)),
-    ]
+    ],
 )
 def test_device_jit(func):
     # Test work around to https://github.com/jax-ml/jax/issues/26000
@@ -36,3 +43,14 @@ def test_device_jit(func):
     x = jnp.ones(1)
     assert_equal(func(x), jnp.asarray([0]))
     assert_equal(jax.jit(func)(x), jnp.asarray([0]))
+
+
+def test_inside_jit():
+    jax = pytest.importorskip("jax")
+    import jax.numpy as jnp
+
+    x = jnp.asarray([1, 2, 3])
+    assert jax.jit(is_jax_array)(x)
+    assert jax.jit(is_array_api_obj)(x)
+    assert not jax.jit(is_writeable_array)(x)
+    assert jax.jit(is_lazy_array)(x)