fix older JAX

Author: crusaderky

array_api_compat/common/_helpers.py

@@ -800,8 +800,11 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
         raise ValueError(f"Unsupported device {device!r}")
     elif is_jax_array(x):
         if not hasattr(x, "__array_namespace__"):
-            # In JAX v0.4.31 and older, this import adds to_device method to x.
+            # In JAX v0.4.31 and older, this import adds to_device method to x...
             import jax.experimental.array_api # noqa: F401
+            # ... but only on eager JAX. It won't work inside jax.jit.
+            if not hasattr(x, "to_device"):
+                return x
         return x.to_device(device, stream=stream)
     elif is_pydata_sparse_array(x) and device == _device(x):
         # Perform trivial check to return the same array if

tests/test_jax.py

@@ -7,14 +7,17 @@
 
 def test_device_jit():
     # Test work around to https://github.com/jax-ml/jax/issues/26000
-    @jax.jit
+    # Also test missing to_device() method in JAX < 0.4.31
+    # when inside jax.jit, even after importing jax.experimental.array_api
+
     def f(x):
         return jnp.zeros(1, device=device(x))
 
-    @jax.jit
     def g(x):
         return to_device(jnp.zeros(1), device(x))
 
     x = jnp.ones(1)
     assert_equal(f(x), jnp.asarray(0))
     assert_equal(g(x), jnp.asarray(0))
+    assert_equal(jax.jit(f)(x), jnp.asarray(0))
+    assert_equal(jax.jit(g)(x), jnp.asarray(0))