Infrastructure / Core ML frameworks & libraries

Flax

Flax is a neural network library and ecosystem for JAX designed for flexibility, providing network components such as Linear, Conv, BatchNorm and attention alongside training utilities like checkpointing and serialization. It is developed and maintained by Google DeepMind. Its newer Flax NNX API emphasizes Python reference semantics, so networks are built from ordinary Python objects.

Verified 2026-08-13 via GitHub and the GitHub license API.

Openness

5 high confidence

5.0

license: Apache-2.0(OSI)
source: public
core-gated: ungated

Apache-2.0 licensed with full public source. The repository is public and unarchived, and the README describes the whole library with no paid, enterprise or hosted tier beside it, so source is public and the core ungated.

https://github.com/google/flax recorded 2026-08-13
Repo page still shows the Apache-2.0 license and the full public source tree
https://api.github.com/repos/google/flax recorded 2026-08-13
Repo metadata - license spdx_id Apache-2.0, private false, archived false, default branch main - for google/flax.
https://raw.githubusercontent.com/google/flax/main/README.md recorded 2026-08-13
README describes Flax as a neural network library and ecosystem for JAX, including the Flax NNX API, installed from PyPI. No paid tier or hosted edition appears.

Adoption

4 high confidence

4.0

4,666,653 PyPI downloads of `flax` in the trailing 30 days, which falls in the 1M-10M a month band of the software usage scale, level 4.

https://pepy.tech/projects/flax recorded 2026-08-13
4,666,653 downloads in the last 30 days for flax

Capability

4 high confidence

4.0

Key NN layer for the JAX ecosystem but built atop JAX rather than foundational itself.

https://github.com/google/flax recorded 2026-08-13
README still describes a neural network library and ecosystem for JAX, including the Flax NNX API

Verified 2026-08-13