a quick patch for JAX 0.7.0 compatibility

Author: rdyro

tests/utils.py

@@ -2,7 +2,7 @@
 import time
 from jax.lib import xla_extension
 
-Device = xla_extension.Device
+Device = jax.Device if hasattr(jax, "Device") else xla_extension.Device
 
 
 def jax_randn(shape, device, dtype):

torch2jax/api.py

@@ -9,7 +9,12 @@
 from torch import Tensor
 import jax
 from jax import ShapeDtypeStruct
-from jax.util import safe_zip
+try:
+    from jax.util import safe_zip
+except ImportError:
+    safe_zip = zip
+
+jax.config.update('jax_use_shardy_partitioner', False)  # TODO: temporary workaround for JAX 0.7.0
 
 # jax version-friendly way of importing the ffi module in jax
 try:

torch2jax/dlpack_passing.py

@@ -12,7 +12,7 @@
 except ImportError:
     from jax.core import ConcretizationTypeError
 
-JAXDevice = jax.lib.xla_extension.Device
+JAXDevice = jax.Device if hasattr(jax, "Device") else jax.lib.xla_extension.Device
 
 
 def _transfer(x: Array | Tensor, via: str = "dlpack", device: str = "cuda"):