Bump optax from 0.2.3 to 0.2.4 (#53)

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]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

requirements.txt

@@ -3,7 +3,7 @@ jax==0.4.30
 jaxlib==0.4.30
 jaxtyping==0.2.34
 chex==0.1.87
-optax==0.2.3
+optax==0.2.4
 tqdm~=4.67
 matplotlib~=3.9
 pdoc~=15.0