Update optax, jax, jaxlib

[cgarciae]

Fixes #176

[openerror]

Confirmed that jax, jaxlib, and optax are updated to the latest versions available through pip.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 87.19%. Comparing base (2558575) to head (cdbfdea).
Report is 40 commits behind head on master.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #177   +/-   ##
=======================================
  Coverage   87.19%   87.19%           
=======================================
  Files         136      136           
  Lines        7433     7433           
=======================================
  Hits         6481     6481           
  Misses        952      952           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.