Future-proof JAX array API import

Author: jakevdp

array_api_compat/common/_helpers.py

@@ -377,9 +377,13 @@ def your_function(x, y):
             elif use_compat is False:
                 import jax.numpy as jnp
             else:
-                # jax.experimental.array_api is already an array namespace. We do
-                # not have a wrapper submodule for it.
-                import jax.experimental.array_api as jnp
+                # JAX v0.4.32 and newer implements the array API directly.
+                # For older JAX versions, it is available via jax.experimental.array_api.
+                import jax.numpy
+                if hasattr(jax.numpy, "__array_api_version__"):
+                    jnp = jax.numpy
+                else:
+                    import jax.experimental.array_api as jnp
             namespaces.add(jnp)
         elif is_pydata_sparse_array(x):
             if use_compat is True: