Package profile
flax
- Summary: Flax: A neural network library for JAX designed for flexibility
- Author: Flax team <flax-dev@google.com>
- Homepage: https://github.com/google/flax
- Source: https://github.com/google/flax (Repo profile)
- Number of releases: 56
- First release: 0.1.0rc1 on 2020-03-17
- Latest release: 0.11.2 on 2025-08-28
Releases
PyPI Downloads
Loading PyPI statistics...
Dependencies
Flax has 51 dependencies, 41 of which optional.| Dependency | Optional |
|---|---|
| jax | false |
| msgpack | false |
| numpy | false |
| optax | false |
| orbax-checkpoint | false |
| PyYAML | false |
| rich | false |
| tensorstore | false |
| treescope | false |
| typing_extensions | false |
| ale-py | true |
| cloudpickle | true |
| clu | true |
| dm-haiku | true |
| docutils | true |
| einops | true |
| gymnasium | true |
| ipykernel | true |
| ipython_genutils | true |
| ipywidgets | true |
| jaxlib | true |
| jaxtyping | true |
| jraph | true |
| jupytext | true |
| kagglehub | true |
| matplotlib | true |
| ml_collections | true |
| ml-collections | true |
| mypy | true |
| myst_nb | true |
| nanobind | true |
| nbstripout | true |
| opencv-python | true |
| pre-commit | true |
| Pygments | true |
| pytest | true |
| pytest-cov | true |
| pytest-custom_exit_code | true |
| pytest-xdist | true |
| pytype | true |
| recommonmark | true |
| scikit-build-core | true |
| scikit-learn | true |
| sentencepiece | true |
| sphinx | true |
| sphinx-book-theme | true |
| sphinx-design | true |
| tensorflow | true |
| tensorflow_datasets | true |
| tensorflow_text | true |
| torch | true |