Flax NNX — The New Standard

Flax NNX is the current recommended API for building neural networks in JAX. It replaced the older Flax Linen API (which required a separate model.init/model.apply dance). NNX feels much more like PyTorch — you create modules as regular Python objects, call them directly, and they hold their own state.

from flax import nnx
import jax.numpy as jnp

# Define a simple model — looks like PyTorch!
class MLP(nnx.Module):
    def __init__(self, in_dim, hidden_dim, out_dim, rngs: nnx.Rngs):
        self.linear1 = nnx.Linear(in_dim, hidden_dim, rngs=rngs)
        self.bn = nnx.BatchNorm(hidden_dim, rngs=rngs)
        self.dropout = nnx.Dropout(rate=0.2, rngs=rngs)
        self.linear2 = nnx.Linear(hidden_dim, out_dim, rngs=rngs)

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

# Create the model — eager initialization, no lazy shapes
model = MLP(784, 256, 10, rngs=nnx.Rngs(0))

# Call it directly — no .apply() needed!
x = jnp.ones((32, 784))
y = model(x)
print(y.shape)  # (32, 10)

Key features of Flax NNX:

• Pythonic: Regular Python objects with mutable state. Debug with print statements, inspect with tab completion.
• Eager initialization: Parameters are created in __init__, not lazily on first call.
• Direct calls: model(x) just works — no model.apply(params, x) ceremony.
• Built-in transforms: nnx.jit, nnx.grad, nnx.vmap handle NNX modules correctly.

# NNX transforms handle modules with state
import optax

model = MLP(784, 256, 10, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3))

# Training step — model and optimizer are mutated in-place
@nnx.jit
def train_step(model, optimizer, x, y):
    def loss_fn(model):
        logits = model(x)
        return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(
            logits, y))

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(grads)
    return loss

# Use it — no return/reassign dance
# loss = train_step(model, optimizer, x_batch, y_batch)

Flax is at version 0.12.x as of early 2026. The main docs at flax.readthedocs.io use NNX by default, while the old Linen API lives at flax-linen.readthedocs.io.