Fixed Optax being a full dependency

ArvinSKushwaha · Jun 8, 2024 · ea32d38 · ea32d38
1 parent 463c425
commit ea32d38
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 22 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,7 +9,6 @@ dependencies = [
     "safetensors>=0.4.3",
     "jaxtyping>=0.2.29",
     "jax>=0.4.28",
-    "optax>=0.2.2",
     "jaxlib>=0.4.28",
 ]
 readme = "README.md"
@@ -22,7 +21,7 @@ build-backend = "hatchling.build"
 
 [tool.rye]
 managed = true
-dev-dependencies = ["pytest>=8.2.1", "basedpyright>=1.12.4"]
+dev-dependencies = ["pytest>=8.2.1", "basedpyright>=1.12.4", "optax>=0.2.2"]
 excluded-dependencies = []
 
 [tool.rye.scripts]

diff --git a/requirements-dev.lock b/requirements-dev.lock
@@ -42,7 +42,6 @@ numpy==1.26.4
 opt-einsum==3.3.0
     # via jax
 optax==0.2.2
-    # via jaximal
 packaging==24.0
     # via pytest
 pluggy==1.5.0

diff --git a/requirements.lock b/requirements.lock
@@ -8,46 +8,27 @@
 #   with-sources: false
 
 -e file:.
-absl-py==2.1.0
-    # via chex
-    # via optax
-chex==0.1.86
-    # via optax
 jax==0.4.28
-    # via chex
     # via jaximal
-    # via optax
 jaxlib==0.4.28
-    # via chex
     # via jaximal
-    # via optax
 jaxtyping==0.2.29
     # via jaximal
 ml-dtypes==0.4.0
     # via jax
     # via jaxlib
 numpy==1.26.4
-    # via chex
     # via jax
     # via jaxlib
     # via ml-dtypes
     # via opt-einsum
-    # via optax
     # via scipy
 opt-einsum==3.3.0
     # via jax
-optax==0.2.2
-    # via jaximal
 safetensors==0.4.3
     # via jaximal
 scipy==1.13.1
     # via jax
     # via jaxlib
-setuptools==70.0.0
-    # via chex
-toolz==0.12.1
-    # via chex
 typeguard==2.13.3
     # via jaxtyping
-typing-extensions==4.12.1
-    # via chex