google/flax

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

How Data Flows Through the System

Data enters through dataset loaders that create batches of examples (images, text tokens, etc.), flows through Module forward passes to produce logits, gets compared against labels in loss functions, generates gradients via automatic differentiation, and updates model parameters through optimizers. The NNX API handles module state automatically while Linen requires explicit parameter passing.

1. Load and preprocess data — Dataset loaders in examples/ read from tensorflow_datasets or custom sources, apply tokenization for text or normalization for images, create batches with proper padding and masking (config: per_device_batch_size, dataset_name, max_corpus_chars)
2. Initialize model and optimizer — Create Module instance with nnx.Rngs for parameter initialization, wrap with Optimizer class that holds optax gradient transformation and tracks parameter state (config: vocab_size)
3. Forward pass through model — Module.__call__ processes input batch through neural network layers (conv, attention, MLP), each layer accesses its Variable parameters and produces intermediate activations [Batch → logits tensor]
4. Compute loss and metrics — Loss functions like cross-entropy compare model logits against target labels, produce scalar loss value plus auxiliary metrics like accuracy for monitoring [logits tensor → loss scalar]
5. Backward pass and parameter updates — nnx.grad automatically differentiates loss with respect to model parameters, optimizer applies gradient updates using momentum, weight decay, learning rate schedules [loss scalar → Updated Module]
6. Checkpoint and evaluation — Training state gets saved to disk every N steps using orbax checkpointing, model runs on validation data to compute eval metrics, logs training progress [Updated Module → saved checkpoint] (config: eval_per_device_batch_size, eval_dataset_name, eval_split)

Data Models

The data structures that flow between stages — the contracts that hold the system together.

Hidden Assumptions

Things this code relies on but never validates. These are the things that cause silent failures when the system changes.

System Behavior

How the system operates at runtime — where data accumulates, what loops, what waits, and what controls what.

Data Pools

Feedback Loops

• Training loop (training-loop, reinforcing) — Trigger: gradient computation. Action: optimizer updates parameters based on loss gradients. Exit: max steps reached or convergence.
• Learning rate scheduling (convergence, balancing) — Trigger: training step count. Action: adjusts learning rate using warmup and decay schedules. Exit: training completion.
• Gradient accumulation (gradient-accumulation, reinforcing) — Trigger: small batch size. Action: accumulates gradients over multiple mini-batches before parameter update. Exit: target effective batch size reached.

Delays

• JIT compilation (compilation, ~seconds on first call) — first training step takes much longer while JAX compiles the computation graph
• Checkpoint saving (checkpoint-save, ~seconds for large models) — training pauses while orbax writes parameter tensors to disk, can be async
• Data loading (batch-window, ~milliseconds per batch) — GPU may idle while waiting for next batch preprocessing and transfer

Control Points

• per_device_batch_size (hyperparameter) — Controls: memory usage and gradient noise during training. Default: 32
• vocab_size (architecture-switch) — Controls: embedding layer dimensions and tokenizer vocabulary. Default: varies by model
• dataset_name (runtime-toggle) — Controls: which dataset pipeline to use for training and evaluation. Default: lm1b
• JAX_PLATFORMS (env-var) — Controls: whether JAX uses CPU, GPU, or TPU devices for computation

Technology Stack

Key Components

• nnx.Module (factory) — Base class for all NNX neural network modules, manages Variable instances and provides graph split/merge functionality for JAX transformations flax/nnx/module.py
• GraphDefState (adapter) — Splits NNX modules into pure structure (GraphDef) and mutable state (State) so they can pass through JAX transformations, then merges them back flax/nnx/graph.py
• Optimizer (optimizer) — Wraps optax optimizers to work with NNX modules, applies gradient updates to module parameters while preserving graph structure flax/nnx/optim.py
• transforms.jit (transformer) — JIT compilation wrapper that automatically handles NNX module graph splitting before compilation and merging after execution flax/nnx/transforms/compilation.py
• transforms.grad (transformer) — Automatic differentiation that works with NNX modules by splitting them before taking gradients and reconstructing gradients in module form flax/nnx/transforms/autodiff.py
• Linen.Module (factory) — Base class for Linen API modules using functional programming style where parameters are managed externally in separate variables dict flax/linen/module.py
• train_state.TrainState (store) — Standard container for training loop state including parameters, optimizer state, step counter, and learning rate schedule flax/training/train_state.py
• checkpoints.save_checkpoint (serializer) — Saves training state to disk using orbax format, handles large model sharding across devices and async writing for performance flax/training/checkpoints.py
• common_utils.shard (dispatcher) — Distributes data and model parameters across multiple devices for data-parallel training, handles device placement and synchronization flax/training/common_utils.py

Related Library Repositories

Frequently Asked Questions

Analyzed on April 20, 2026 by CodeSea. Written by Karolina Sarna.