From 62517feb940a1a0c255d0bbaa33c2b8c9c33b64a Mon Sep 17 00:00:00 2001
From: Jeffrey Gerard <IamJeffG@users.noreply.github.com>
Date: Mon, 5 Feb 2024 20:26:47 -0600
Subject: [PATCH] Constrain optax to 0.1.x

See https://github.com/google-research/perch/issues/613
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 = [