Flax is Google’s neural-network library built directly on top of JAX, and it occupies a peculiar position in the deep-learning ecosystem: it is the closest thing the functional-programming world has to a household framework. Where PyTorch gives you a mutable object whose weights live in place, Flax — at least in its original Linen API — insists that a model is a pure function and that its weights are data you pass in and out explicitly. That single design decision explains almost everything else in the codebase: the FrozenDict containers, the init/apply duality, the TrainState bundle, and eventually the decision to build an entirely second API, NNX, on top of the same substrate.

The library is also a survivor. It powered some of the largest open training efforts in the JAX world — big-scale projects in the MaxText family still follow Flax-style patterns — and yet by 2024 the team had concluded that Linen’s ceremony was a tax on everyday users. The result is flax.nnx, a redesign where modules are ordinary mutable Python objects again, with the functional machinery moved behind explicit split/merge calls. Today the repository carries both APIs side by side, and reading the source means reading two philosophies of state management implemented against one compiler.

This post is a source tour of google/flax on main. We walk the Linen module system (nn.Module, setup versus @nn.compact), the transform system (nn.vmap, nn.scan), the parameter containers, TrainState, and then the NNX graph machinery that makes mutable modules composable with JAX transforms. The goal is that after this tour, neither API feels like magic.

The two halves of Flax: Linen (functional modules plus lifts and state containers on the left) and NNX (mutable module graphs with explicit split/merge on the right), meeting at shared JAX pytree infrastructure.

Reading the overview from left to right:

• nn.Module, L979 is the Linen base class. You declare submodules either imperatively in setup, L1391 or inline inside a method decorated with @nn.compact, L476. The class itself never stores parameters — it is a shape of a computation, not its weights.
• nn.vmap / nn.scan, L834 and its sibling scan, L1153 are the “lifts”: wrappers that take a module method and re-execute it under JAX’s vmap or scan while automatically threading parameter state in and out.
• nn.tabulate, L174 is the introspection tool. It traces an apply call and renders the module tree, activation shapes, and per-submodule parameter counts as a text table.
• FrozenDict, L54 is the immutable parameter container that init returns. Its unfreeze, L159 method hands you a plain mutable dict, which is also how gradients get applied.
• TrainState, L25 bundles step, params, tx, and apply_fn into one pytree so the entire training state can be jax.jit-compiled and passed around as a value. It is built on PyTreeNode, L194, the frozen-dataclass helper that registers classes as JAX pytrees.
• nnx.Module, L55 restarts the story: submodules and variables are ordinary attributes you assign in __init__, and the object is mutable in place.
• split / merge, L2268 is NNX’s escape hatch into functional land — merge, L2395 is its inverse — and nnx.TrainState, L300 wraps the pattern for training loops.
• Rngs, L343 manages named random-key streams; unlike Linen’s ad-hoc make_rng, an Rngs object is itself a variable that lives on the module.

Two APIs, One Library

The most striking thing about the Flax repository is that it contains two complete module systems. Linen is the mature one: every module is a dataclass-like descriptor and all state flows through explicit function arguments. NNX is the redesign: state lives on the object, and the functional view is something you opt into per transform.

In Linen, a dense layer looks like a factory. Inside @nn.compact, nn.Dense, L214 calls self.param to declare its kernel, and the first time the module runs, that call materializes weights from an initializer and a PRNG key. The key comes from make_rng, L1912, which pops a stream named params out of the scope the caller provided. Nothing is stored: run init and you get a FrozenDict of weights; run apply with those weights and you get outputs.

NNX flips the polarity. nnx.Linear, L313 assigns self.kernel = nnx.Param(...) in __init__, and that Variable, L1098 object — whose Param, L2035 subclass marks it as trainable — stays on the instance. You call the model like a plain Python object. If you have used PyTorch, this feels immediately familiar; the difference is that every NNX transform knows how to extract the variables, do the pure functional thing underneath, and write the results back.

The bridge between the two worlds is the graph machinery. Because NNX modules form a mutable object graph, JAX cannot trace them directly. nnx.split walks the graph and flattens it into a GraphDef (static topology) plus one or more State, L221 mappings (the actual arrays); merge rebuilds a live module from those pieces. Filters let you split only some variable types — nnx.split(model, nnx.Param) for parameters, a different filter for batch statistics — which is how NNX reproduces the collection-based separation Linen did with named collections.

From Install to Your First Model

Flax installs from PyPI with pip install flax. It pulls in JAX automatically; whether you get GPU or TPU acceleration depends on the JAX backend you install alongside it, exactly as with any JAX program. There is no compiled extension to build and no system dependency to hunt down.

Here is a complete NNX model, taken almost verbatim from the docstring of the module class itself:

import jax.numpy as jnp
from flax import nnx

class Model(nnx.Module):
    def __init__(self, rngs):
        self.linear1 = nnx.Linear(2, 3, rngs=rngs)
        self.linear2 = nnx.Linear(3, 4, rngs=rngs)
    def __call__(self, x):
        x = self.linear1(x)
        x = nnx.relu(x)
        return self.linear2(x)

model = Model(rngs=nnx.Rngs(0))
y = model(jnp.ones((1, 2)))

Every constructor that creates variables takes rngs= — an Rngs object that hands out named key streams. From there you can drop into the functional view whenever you need it:

graphdef, params = nnx.split(model, nnx.Param)
# ... inspect or serialize `params` as a pytree ...
model = nnx.merge(graphdef, params)

And the training-loop idiom is to let nnx.jit (flax/nnx/transforms/compilation.py, L200) handle the split/merge sandwich for you: you write an ordinary function that mutates the model, decorate it with @nnx.jit, and the transform extracts state before the call and applies updates after it. For comparison, the same model in Linen requires an init call to produce a parameter tree and an apply call for every forward pass — apply, L2090 takes the params as its second argument, while init, L2314 creates them from an RNG key. One API treats state as a detail; the other makes it the interface.

The Functional Core: init, apply, and FrozenDict

Linen’s design is easiest to understand as a discipline imposed on JAX. JAX’s jit and grad want pure functions over pytrees, so Flax makes the model a pure function and the parameters a pytree. The pytree in question is FrozenDict, a hash-frozen, immutable mapping. Immutability is not aesthetic: it guarantees that a traced computation cannot silently mutate the weights it was compiled against, and it makes parameter trees safe to share across threads and cache across jit calls.

To update parameters you call unfreeze, apply your gradient, and re-wrap — or, more commonly, you let TrainState do it. Its apply_gradients, L81 method takes gradients and returns a new TrainState with updated params and an incremented step counter, and its create, L125 classmethod wires up the optimizer from optax. Because the whole object is a pytree (courtesy of PyTreeNode), one jax.jit covers the step function. This pattern spread far beyond Flax — plenty of JAX codebases that never adopted Linen still have a TrainState-shaped object at their center.

There is a subtlety worth knowing: parameters are not the only state. self.param declares read-only variables in the params collection, while the more general variable, L1720 method can declare mutable collections — batch_stats for normalization running means, cache for autoregressive decoding, anything you name. apply accepts mutable=[...] to say which collections may change, and returns them alongside the output. NNX generalizes this: every variable carries its type, and filters select by type, so the collection idea became a class hierarchy.

Lifts: vmap, scan, and Friends

The transforms are where Flax earns its complexity budget. nn.vmap maps a module method over a batch axis, but unlike raw jax.vmap it also maps the parameters: each logical “member” of the mapped module gets its own slice of the parameter tree, sliced automatically along the axes you nominate. nn.scan does the same for a loop over time, threading carry state through iterations and stacking the per-step outputs.

Implementing this on immutable state is genuinely intricate — the lift has to track which variables are mapped, which are carried, and which are broadcast, then reassemble the state after the call — which is why the file runs well over a thousand lines and sits on flax.core.lift, the lower-level lifting machinery. On top of it, scan_with_axes, L416 and vmap_with_axes, L472 in the partitioning module add support for AxisMetadata, so that sharding annotations survive the lift — the detail that made large-scale MaxText-style training practical in Flax.

NNX re-derived all of this for the object world. nnx.vmap, L398 and nnx.scan, L1499 split the module, perform the lifted computation on pytrees, and merge the mutated state back into the live objects afterward. The supporting helpers split_inputs, L43 and merge_inputs, L172 convert module arguments into jax-compatible data so that arbitrary transforms can be applied without special-casing every call site. The user-facing result is nicer — you decorate a method, not a class — but the source shows the same conceptual machinery underneath, just relocated from the caller’s hands into the library’s.

TrainState and the Training Loop

A canonical Linen training step is almost boring in its regularity: build a TrainState once with create, then loop grads = grad(loss)(state.params), state = state.apply_gradients(grads=grads), all under one jit. The step counter, optimizer, and parameters travel together, which makes checkpointing trivial — the entire training state is one pytree you can serialize, and libraries in the wider JAX orbit (including the Hugging Face transformers JAX integrations in their day) adopted the shape.

NNX’s nnx.TrainState mirrors the idea, but because the module is mutable, the class is generic over the module type and uses split/merge internally to expose a functional surface. Notably, its source imports the struct helper from flax.training.train_state — the two APIs share their pytree plumbing even where their surfaces diverge. If you are porting a Linen codebase, this is the pattern to lean on: keep the TrainState-shaped loop, swap the module construction and the apply call, and let the graph machinery handle the stateful parts.

One practical wrinkle: in NNX the “which variables are trainable” question is answered by filters at the split site, so the same module can be wrapped differently for training versus evaluation — parameters only, or parameters plus batch statistics. It is flexible, but it moves a decision that Linen encoded in collection names into your call sites.

Honest Limits

The dual-API situation is the elephant. If you read Flax documentation today, you will find NNX presented as the recommended path while enormous amounts of production code — and many published code examples, checkpoints, and university course materials — are Linen. The two APIs are close enough in spirit that translation is mechanical, but not drop-in: collection strings become variable-type filters, setup becomes __init__, and the mutable= argument becomes filter lists. Teams carry both mental models, and the repository carries both test suites and both summary tools (nn.tabulate for Linen, the object-level display machinery for NNX).

The functional ceremony is the second real cost. In Linen, a beginner’s first training loop involves concepts PyTorch never asks for: separate init and apply, immutable dictionaries, named RNG streams, and the mutable= flag. None of it is gratuitous — each piece exists because JAX needs purity — but the on-ramp is objectively steeper, and error messages when you pass the wrong state shape can be opaque. NNX fixes much of the ergonomics at the price of a new set of rules about when state is extracted and re-applied, and transform interactions (scan over a module that jits another module) still require care in either API.

Finally, the docs lag the code in places. The NNX API has moved fast — module, variable, and state internals were reworked repeatedly — so blog posts and Stack Overflow answers often describe stale signatures. Check the source before trusting a snippet: the repository is well-commented, docstrings carry doctests, and as this tour shows, the whole system is small enough to read in an afternoon. If you want the compiler-side context for why all this functional machinery exists, our JAX tour is the natural next stop; for a different take on the same mutable-object problem, Keras shows what a framework looks like when it hides the functional layer entirely.

Enjoyed this post? Never miss out on future posts by following us