# google/flax

Flax is a neural network library for JAX that is designed for flexibility.

Repository: https://github.com/google/flax
Canonical: https://ross.abutalabs.com/products/flax
Homepage: https://flax.readthedocs.io
Language: Jupyter Notebook
License: Apache-2.0
License Family: permissive
Topics: jax
Last push: 2026-08-24T19:34:32+00:00

## Health v2 (maintenance only)
Score: 99/100 (v2, computed 2026-09-03T02:20:16.233290+00:00)
- activity 99, release rhythm 98, longevity 100
- inputs: {"age_days": 2427, "days_push": 9, "days_rel": 15, "gap_med": 27.5, "n_releases_24m": 21}
- flags: none
- formula: round(0.45*activity + 0.35*rhythm + 0.20*longevity); archived -> min(score, 10)

## Adoption (not part of the score)
Stars 7303, forks 830 (observed 2026-08-28T04:09:58.132098+00:00)

## What it is
Flax is a neural network library and ecosystem for JAX designed for flexibility, featuring the newer NNX API with first-class Python reference semantics alongside the older Linen API. It is developed in close collaboration with the JAX team for high-performance deep learning research.

## Use cases
- build neural networks in jax
- train deep learning models with jax
- define models with regular python objects and mutable state
- run deep learning research with flexible training loops
- migrate from flax linen to the nnx api
- train models on gpus and tpus

## When to choose
- you use JAX and want a flexible, research-friendly neural network library
- you want Python reference semantics with mutable, shareable model state (NNX)
- you need a mature, Google-backed ecosystem with examples and tutorials

## When to avoid
- you prefer PyTorch or TensorFlow-style frameworks
- you want a batteries-included high-level training framework with opinionated defaults
- your project depends on the legacy Linen API without plans to migrate

## Facets
- artifact type: library
- maturity: active
- function: machine-learning, deep-learning, llm-training
- domain: machine-learning, deep-learning, artificial-intelligence
- platform: python, cross-platform
- tags: jax, neural-networks, nnx, linen, research, google, gpu

## Member repositories
- google/flax (main) score 99

## Provenance
- Observed fields: from GitHub, fetched 2026-08-28T04:09:58.132098+00:00.
- Health v2: computed from the inputs above; adoption is never an input.
- Inferred fields (summary, facets, guidance): AI-extracted, prompt v1, taxonomy v1, on 2026-08-29T17:38:41.740331+00:00, confidence not recorded.
  - readme: https://github.com/google/flax (fetched 2026-08-28T04:09:58.132098+00:00, sha 9c7d9f7f40c4)
  - registry_pypi: https://pypi.org/pypi/flax/json (fetched 2026-08-29T08:34:33.506554+00:00, sha 94ff4bfef308)
- Data as of 2026-08-30T08:39:29.467469+00:00.