diff --git a/setup.py b/setup.py
index f7d6c12d..dbcf438b 100644
--- a/setup.py
+++ b/setup.py
@@ -80,7 +80,7 @@ def run(self):
     "immutabledict>=2.2.1",
     "clu>=0.0.6",
     "tensorflow-datasets",
-    "optax @ git+https://github.com/google-deepmind/optax.git@main",
+    "optax>=0.1.9,==0.1.*",
 ]
 
 tests_require = [