Package profile
flax
- Summary: Flax: A neural network library for JAX designed for flexibility
- Author: Flax team <flax-dev@google.com>
- License: Apache Software License
- Homepage: https://github.com/google/flax
- Source: https://github.com/google/flax
- Number of releases: 67
- First release: 0.1.0rc1 on 2020-03-17
- Latest release: 0.12.10 on 2026-09-24
- Latest release size: 519.7 KB (pure Python wheel)
Dependencies
Flax has 56 dependencies, 46 of which optional.Dependent packages
| Package | Optional | Group |
|---|---|---|
| brax | false | |
| clu | false | |
| dopamine-rl | false | |
| evojax | false | |
| evosax | false |
Similar packages
- optaxA gradient processing and optimization library in JAX.
- cluSet of libraries for ML training loops in JAX.
- mxnetApache MXNet is an ultra-scalable deep learning framework. This version uses openblas and MKLDNN.
- jaxlibXLA library for JAX
- onnxruntimeONNX Runtime is a runtime accelerator for Machine Learning models
- onnxruntime-gpuONNX Runtime is a runtime accelerator for Machine Learning models
- distraxDistrax: Probability distributions in JAX.