Back to Gap Map Infrastructure / Core ML frameworks & libraries

Flax

open source / Overall score: 4.0(strong)

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.

Openness

5 high confidence

5.0

license: Apache-2.0(OSI)
source: public
core features withheld: no

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://cdn.jsdelivr.net/gh/google/flax@main/LICENSE recorded 2026-09-29
The LICENSE file is the full text of the "Apache License" "Version 2.0, January 2004", so the code is Apache-2.0 licensed.
https://cdn.jsdelivr.net/gh/google/flax@main/README.md recorded 2026-09-29
The README describes Flax as "a high-performance neural network library and ecosystem for" JAX. Searched the body for paid, enterprise, hosted, license key and self-host: none, so no gated tier sits beside the library.
https://flax.readthedocs.io/en/latest/contributing.html recorded 2026-09-29
The contributing guide tells readers to "fork the Flax repository by clicking the" fork button on GitHub, so the source is public.
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://pypi.org/pypi/flax/json recorded 2026-09-29
PyPI metadata for flax lists the classifier "License :: OSI Approved :: Apache Software License" and the homepage "https://github.com/google/flax".

Adoption

4 high confidence

4.0

Adoption is measured by PyPI downloads of the flax package.

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://cdn.jsdelivr.net/gh/google/flax@main/README.md recorded 2026-09-29
The README lists the Flax NNX neural network API with Linear, Conv, BatchNorm, attention and recurrent cells, plus "replicated training, serialization and checkpointing, metrics, prefetching on device".
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-09-24