Package profile
optax
- Summary: A gradient processing and optimization library in JAX.
- Author: Google DeepMind <optax-dev@google.com>
- License: Apache Software License
- Homepage: https://github.com/google-deepmind/optax
- Documentation: https://optax.readthedocs.io/
- Source: https://github.com/google-deepmind/optax
- Number of releases: 26
- First release: 0.0.1 on 2020-08-14
- Latest release: 0.2.8 on 2026-03-20
- Latest release size: 393.5 KB (pure Python wheel)
Dependencies
Optax has 18 dependencies, 14 of which optional.Dependent packages
| Package | Optional | Group |
|---|---|---|
| brax | false | |
| flax | false | |
| flowjax | false | |
| nequix | false | |
| vizopt | false |
Similar packages
- flaxFlax: A neural network library for JAX designed for flexibility
- cluSet of libraries for ML training loops in JAX.
- distraxDistrax: Probability distributions in JAX.
- optunaA hyperparameter optimization framework
- jaxlibXLA library for JAX
- nevergradA Python toolbox for performing gradient-free optimization
- diffusersState-of-the-art diffusion in PyTorch and JAX.