Merge branch 'master' of github.com:wesselb/lab

wesselb · Mar 6, 2023 · 51b1b0a · 51b1b0a
2 parents c0da7de + 7070394
commit 51b1b0a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lab/types.py b/lab/types.py
@@ -86,7 +86,10 @@ def _module_attr(module, attr):
 _ag_tensor = ModuleType("autograd.tracer", "Box")
 
 # Define JAX module types.
-_jax_tensor = ModuleType("jax.interpreters.xla", "DeviceArray")
+if sys.version_info.minor > 7:  # jax>0.4 deprecated python-3.7 support, rely on older jax versions
+    _jax_tensor = ModuleType("jaxlib.xla_extension", "ArrayImpl")
+else:
+    _jax_tensor = ModuleType("jax.interpreters.xla", "DeviceArray")
 _jax_tracer = ModuleType("jax.core", "Tracer")
 _jax_dtype = ModuleType("jax._src.numpy.lax_numpy", "_ScalarMeta")
 _jax_device = ModuleType("jaxlib.xla_extension", "Device")