fix: Pin orbax-checkpoint to v0.10.3 to resolve dependency error (AI-…

…Hypercomputer#1273) - Update requirements.txt and requirements_with_jax_stable_stack.txt to specify orbax-checkpoint==0.10.3. - Prevent AttributeError in MaxText/train.py related to key_value_try_get.
shota-inoue-lts · Feb 14, 2025 · 58e9045 · 58e9045
1 parent 39b5b74
commit 58e9045
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,6 @@
 jax>=0.4.30
 jaxlib>=0.4.30
-orbax-checkpoint>=0.5.12
+orbax-checkpoint==0.10.3
 absl-py
 array-record
 aqtp

diff --git a/requirements_with_jax_stable_stack.txt b/requirements_with_jax_stable_stack.txt
@@ -4,7 +4,7 @@ absl-py
 aqtp==0.8.2
 datasets
 grain-nightly>=0.0.10
-orbax-checkpoint>=0.10.3
+orbax-checkpoint==0.10.3
 pylint
 pytest
 pyink