Add comments on jax.core.Tracer detection limitations

Author: amacati

array_api_compat/common/_helpers.py

@@ -235,6 +235,12 @@ def is_jax_array(x: object) -> TypeIs[jax.Array]:
     is_pydata_sparse_array
     """
     cls = cast(Hashable, type(x))
+    # We test for jax.core.Tracer here to identify jax arrays during jit tracing. From jax 0.8.2 on,
+    # tracers are not a subclass of jax.Array anymore. Note that tracers can also represent
+    # non-array values and a fully correct implementation would need to use isinstance checks. Since
+    # we use hash-based caching with type names as keys, we cannot use instance checks without
+    # losing performance here. For more information, see
+    # https://github.com/data-apis/array-api-compat/pull/369 and the corresponding issue.
     return (
         _issubclass_fast(cls, "jax", "Array")
         or _issubclass_fast(cls, "jax.core", "Tracer")
@@ -300,7 +306,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")
+        or _issubclass_fast(cls, "jax.core", "Tracer")  # see is_jax_array for limitations
     )
 
 
@@ -939,7 +945,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, "jax.core", "Tracer")  # see is_jax_array for limitations
         or _issubclass_fast(cls, "sparse", "SparseArray")
     ):
         return False
@@ -979,7 +985,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, "jax.core", "Tracer")  # see is_jax_array for limitations
         or _issubclass_fast(cls, "dask.array", "Array")
         or _issubclass_fast(cls, "ndonnx", "Array")
     ):

tests/test_jax.py

@@ -46,9 +46,10 @@ def test_device_jit(func):
 
 
 def test_inside_jit():
-    jax = pytest.importorskip("jax")
-    import jax.numpy as jnp
-
+    # Test if jax arrays are handled correctly inside jax.jit.
+    # Jax tracers are not a subclass of jax.Array from 0.8.2 on. We explicitly test that
+    # tracers are handled appropriately. For limitations, see is_jax_array() docstring.
+    # Reference issue: https://github.com/data-apis/array-api-compat/issues/368
     x = jnp.asarray([1, 2, 3])
     assert jax.jit(is_jax_array)(x)
     assert jax.jit(is_array_api_obj)(x)