Pin optax<0.2.7 to avoid tree flatten_up_to regression

optax 0.2.8 has a bug where its optimizer state contains None entries
that jax.tree.map's flatten_up_to no longer accepts as tree prefixes,
breaking gradient updates in test_hf_gpt2_serialize.

Author: rjpower

lib/levanter/pyproject.toml

@@ -39,7 +39,7 @@ dependencies = [
     "tokenizers>=0.15.2",
     "transformers>=4.57.1,<5.0",
     "chex>=0.1.86",
-    "optax>=0.1.9",
+    "optax>=0.1.9,<0.2.7",
     "wandb>=0.17.8",
     "draccus>=0.11.5",
     "pyarrow>=11.0.0",