Infrastructure / Core ML frameworks & libraries

Flax

Flax is a neural network library and ecosystem for JAX designed for flexibility, providing network components (Linear, Conv, BatchNorm, attention) plus training utilities like checkpointing and serialization. It is developed and maintained by Google DeepMind as an open source project. Its newer Flax NNX API emphasizes Python reference semantics for building networks with standard Python objects.

Verified live 2026-06-22 via primary sources. Apache-2.0 licensed with full public source.

Openness

5 high confidence

5.0

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

Apache-2.0 licensed with full public source.

https://github.com/google/flax recorded 2026-06-22
Repo displays Apache-2.0 license and full source

Adoption

4 high confidence

4.0

5.3M PyPI downloads in the last 30 days, just above the 5M Level-4 threshold.

https://pepy.tech/projects/flax recorded 2026-06-22
5,302,085 downloads in last 30 days

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-06-22
README describes neural network library and ecosystem for JAX

Unchanged since 2026-07-30 (last edited, not re-checked)