pin lowest jax dependencies

Author: pattonw

examples/jax/identity.py

@@ -69,7 +69,7 @@ def __call__(self, x):
 ckpt_path = work_dir / "ckpt"
 checkpointer = ocp.StandardCheckpointer()
 checkpointer.save(str(ckpt_path.resolve()), params)
-checkpointer.wait_until_finished()
+checkpointer.close()
 
 # %% [markdown]
 # ## 3. Configure JaxModel and Predict

pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
 
 [project.optional-dependencies]
 torch = ["torch>=2.4"]
-jax = ["jax>=0.4", "flax>=0.8", "orbax-checkpoint>=0.5"]
+jax = ["jax>=0.4.30", "flax>=0.8", "orbax-checkpoint>=0.5"]
 dev = [
     "pytest>=8.0",
     "pytest-cov>=5.0",

tests/test_models.py

@@ -110,7 +110,7 @@ def test_jax_models(save_type, out_range, tmp_path):
         params_path = tmp_path / "ckpt"
         checkpointer = ocp.StandardCheckpointer()
         checkpointer.save(str(params_path), params)
-        checkpointer.wait_until_finished()
+        checkpointer.close()
     elif save_type == "msgpack":
         from flax.serialization import to_bytes