Pinning JAX to <0.4.31 until sharding parameter bug is solved.

graphcore-research · Aug 12, 2024 · 759e9e5 · 759e9e5
1 parent 9f00286
commit 759e9e5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,11 +23,11 @@ classifiers = [
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dependencies = [
- "chex >= 0.1.6",
- "jax >= 0.3.16",
- "jaxlib >= 0.3.15",
+ "chex>=0.1.6",
+ "jax>=0.3.16,<0.4.31",
+ "jaxlib>=0.3.15",
  "ml_dtypes",
- "numpy >= 1.22.4"
+ "numpy>=1.22.4"
 ]
 dynamic = ["version"]