Bump optax from 0.2.3 to 0.2.4

Bumps [optax](https://github.com/google-deepmind/optax) from 0.2.3 to 0.2.4.
- [Release notes](https://github.com/google-deepmind/optax/releases)
- [Commits](google-deepmind/optax@v0.2.3...v0.2.4)

---
updated-dependencies:
- dependency-name: optax
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <[email protected]>

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "jax==0.4.35",
     "flax==0.10.1",
     "ml_dtypes==0.4.0",
-    "optax==0.2.3",
+    "optax==0.2.4",
     "orbax-checkpoint==0.8.0",
     "orbax-export==0.0.5",
 ]