Forcing a modern version of JAX.

GridTools · Oct 2, 2024 · e2aa76f · e2aa76f
1 parent 951f857
commit e2aa76f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -25,7 +25,7 @@ classifiers = [
 ]
 dependencies = [
   "dace>=0.16",
-  "jax[cpu]>=0.4.24",
+  "jax[cpu]>=0.4.33",
   "numpy>=1.26.0",
 ]
 description = "JAX jit using DaCe (Data Centric Parallel Programming)"

diff --git a/src/jace/__init__.py b/src/jace/__init__.py
@@ -9,12 +9,18 @@
 
 from __future__ import annotations
 
+import jax
+
 import jace.translator.primitive_translators as _  # noqa: F401 [unused-import]  # Needed to populate the internal translator registry.
 
 from .__about__ import __author__, __copyright__, __license__, __version__, __version_info__
 from .api import grad, jacfwd, jacrev, jit
 
 
+if jax.version._version_as_tuple(jax.__version__) < (0, 4, 33):
+    raise ImportError(f"Require at least JAX version '0.4.33', but found '{jax.__version__}'.")
+
+
 __all__ = [
     "__author__",
     "__copyright__",