google/flax

What is Flax?

Flax is an open‑source neural‑network library built on top of JAX, a high‑performance numerical computing framework. It provides a flexible, Python‑friendly API (the newer Flax NNX API) that lets researchers write models as regular Python objects, with support for reference sharing and mutability. The library includes common layers (Linear, Conv, BatchNorm, Attention, LSTM/GRU, Dropout, etc.), utilities for replicated training, checkpointing, and metrics, and a collection of educational examples such as MNIST and a Gemma language‑model demo.

Who maintains it?

Developed by engineers and researchers at Google DeepMind in close collaboration with the JAX team. It is not an official Google product but is actively maintained and open to community contributions via GitHub discussions and pull requests.

How do you get started?

1. Install JAX (follow the JAX CPU/GPU/TPU guide).
2. Install Flax via PyPI:

   pip install flax
   

   Optional: pip install "flax[all]" for extra dependencies like Matplotlib.
3. Write a model by subclassing nnx.Module and using the provided layers, then train it with standard JAX code.

Example code (from the README)

class MLP(nnx.Module):
  def __init__(self, din, dmid, dout, *, rngs):
    self.linear1 = nnx.Linear(din, dmid, rngs=rngs)
    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
    self.bn = nnx.BatchNorm(dmid, rngs=rngs)
    self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)

  def __call__(self, x):
    x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))
    return self.linear2(x)

Where to learn more?

• Documentation site: https://flax.readthedocs.io/
• Tutorials: MNIST tutorial, Gemma LM inference example, and the "Flax NNX basics" guide.
• Discussions & support: GitHub Discussions, issue tracker, and the flax-dev@google.com mailing address.

When to use Flax?

If you are already using JAX and need a neural‑network library that:

• Keeps the full flexibility of JAX (no hidden graph compilation steps).
• Allows you to write models with ordinary Python semantics.
• Provides ready‑made layers, training utilities, and example code. then Flax is a natural choice.

---

Citation: The README provides a BibTeX entry for academic referencing.