Add optax dependency

thebuckleylab · Jun 18, 2024 · 4e31fff · 4e31fff
1 parent eefc2ba
commit 4e31fff
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -19,7 +19,8 @@ pip install jpc
 ```
 
 Requires Python 3.9+, JAX 0.4.23+, [Equinox](https://github.com/patrick-kidger/equinox) 
-0.11.2+, [Diffrax](https://github.com/patrick-kidger/diffrax) 0.5.1+, and 
+0.11.2+, [Diffrax](https://github.com/patrick-kidger/diffrax) 0.5.1+, 
+[Optax](https://github.com/google-deepmind/optax) 0.2.2+, and 
 [Jaxtyping](https://github.com/patrick-kidger/jaxtyping) 0.2.24+.
 
 ## Documentation

diff --git a/docs/index.md b/docs/index.md
@@ -18,7 +18,8 @@ pip install jpc
 ```
 
 Requires Python 3.9+, JAX 0.4.23+, [Equinox](https://github.com/patrick-kidger/equinox) 
-0.11.2+, [Diffrax](https://github.com/patrick-kidger/diffrax) 0.5.1+, and 
+0.11.2+, [Diffrax](https://github.com/patrick-kidger/diffrax) 0.5.1+, 
+[Optax](https://github.com/google-deepmind/optax) 0.2.2+, and 
 [Jaxtyping](https://github.com/patrick-kidger/jaxtyping) 0.2.24+.
 
 ## Quick example

diff --git a/pyproject.toml b/pyproject.toml
@@ -36,6 +36,7 @@ dependencies = [
     "jax>=0.4.23",
     "equinox>=0.11.2",
     "diffrax>=0.5.1",
+    "optax>=0.2.2",
     "jaxtyping>=0.2.24"
 ]