Flax

Flax Linen provides a functional programming model for building neural networks where modules are defined as classes inheriting from flax.linen.Module, with explicit separation of parameters (immutable) and state through the Scope system. The framework uses a two-phase initialization pattern: init() creates parameters via JAX transformations, and apply() executes forward passes with frozen parameters, eliminating hidden state mutations and enabling seamless composition with JAX's jit, vmap, and grad transformations. State is managed through flax.core.scope.Scope objects that track variable collections (params, batch_stats, cache) hierarchically.

Flax NNX (Neural Network eXperimental) provides a Python-native, object-oriented API released in 2024 where modules are regular Python classes with mutable attributes representing parameters, state, and buffers. The framework uses a GraphDef/State splitting pattern (flax/nnx/graph.py) that separates static module structure from dynamic values, enabling JAX transformations to work with stateful objects. Variables are tracked through flax.nnx.variablelib.Variable subclasses (Param, BatchStat, Cache) that are automatically discovered via Python's attribute introspection, eliminating the need for explicit Scope management while maintaining functional purity during transformations.

Flax provides module lifecycle hooks (setup(), __call__(), __post_init__() for NNX) that enable custom layer implementations with explicit variable creation and management. In Linen, setup() is called once during initialization to create parameters, while __call__() defines the forward pass; in NNX, __post_init__() initializes mutable attributes and __call__() executes forward logic. The framework automatically discovers variables through attribute introspection (NNX) or explicit variable creation within Scope (Linen), enabling custom layers to integrate seamlessly with Flax's variable system, transformations, and checkpointing without manual state threading.

Flax models are represented as PyTrees (nested dicts/lists of JAX arrays) that can be serialized using standard Python libraries (pickle, msgpack, safetensors) or Orbax's checkpoint format. The framework provides utilities for converting Flax models to inference-optimized formats, including parameter quantization, pruning, and conversion to ONNX or TensorFlow SavedModel for cross-framework deployment. PyTree structure enables efficient serialization without framework-specific overhead, and Flax provides helpers for loading models in inference-only mode without optimizer state.

Implements functional random number generation using JAX's PRNG key system, where randomness is explicit and reproducible through key splitting (jax.random.fold_in, jax.random.split). Flax modules use dropout_rng and other random collections to manage randomness during training, with keys automatically split across layers and timesteps. This enables deterministic training with explicit control over randomness, unlike PyTorch's global random state.

Flax provides lifted versions of JAX's core transformations (jit, vmap, scan, pmap) through flax.linen.transforms and flax.nnx.transforms that automatically handle variable state during transformation application. These lifted transforms use a variable collection system where parameters are frozen (non-transformed), while mutable collections like batch_stats and cache are properly threaded through transformation boundaries. For example, nn.vmap automatically batches over specified axes while keeping parameters shared, and nn.scan unrolls recurrent operations while managing state updates, eliminating the need for manual state threading that would be required with raw JAX transformations.

Flax provides flax.training.train_state.TrainState, a dataclass that bundles model parameters, optimizer state, and training metadata (step count, learning rate schedule) into a single immutable structure. TrainState integrates with Optax optimizers through a standard apply_gradients() pattern that atomically updates parameters and optimizer state in a single functional operation. The structure is designed for seamless checkpointing with Orbax (flax/training/checkpoints.py), enabling save/restore of complete training state including optimizer momentum, learning rate schedules, and custom metrics without manual serialization logic.

Flax integrates with Orbax (Google's checkpoint library) through flax/training/checkpoints.py to provide distributed-aware checkpoint save/restore with automatic sharding, async I/O, and incremental updates. The integration handles PyTree serialization of TrainState and model parameters, automatically managing distributed checkpoints across multiple hosts/devices without requiring manual synchronization logic. Orbax's CheckpointManager handles versioning, cleanup of old checkpoints, and recovery from partial writes, while Flax's wrapper provides convenience functions for common patterns like periodic checkpointing during training.

Flax provides a comprehensive library of neural network layers (Dense, Conv2D, LSTM, Attention, Normalization, etc.) in flax.linen.nn and flax.nnx.nn, each implemented with JAX-specific optimizations and variable management. Layers are designed as composable modules that work seamlessly with both Linen's functional API and NNX's OOP API. The library includes architecture-specific implementations like multi-head attention (flax.linen.MultiHeadDotProductAttention) with optional caching for efficient inference, batch normalization with configurable momentum, and dropout with proper PRNG handling, all integrated with Flax's variable collection system.

Flax provides SPMD (Single Program Multiple Data) parallelism support through flax.linen.transforms.pmap and flax.nnx.transforms.pmap, which automatically handle variable sharding across devices/hosts using JAX's pmap primitive. The framework uses axis annotations (via flax.linen.partitioning or manual axis specifications) to declare which dimensions should be parallelized, and automatically threads sharded variables through the computation graph. For distributed training, Flax integrates with JAX's collective operations (all-reduce, all-gather) to synchronize gradients across devices, with built-in support for gradient accumulation and loss scaling for mixed-precision training.

Flax provides flax.linen.summary.summary() and flax.linen.summary.tabulate() functions that introspect module structure to generate human-readable summaries of model architecture, parameter counts, and computational complexity. These tools use Flax's module lifecycle hooks (setup(), __call__()) to trace module instantiation and generate parameter tables showing layer names, shapes, and counts. The summary system works by running a dry-pass through the module with abstract JAX arrays, capturing variable creation without actual computation, enabling fast architecture visualization without GPU memory requirements.

Flax provides flax.linen.dtypes module for managing numerical precision across models, including automatic mixed-precision (AMP) support through dtype specifications on layers and modules. The framework allows per-layer dtype configuration (float32, float16, bfloat16) with automatic casting of inputs/outputs, and supports loss scaling for stable mixed-precision training. Flax integrates with JAX's dtype promotion rules to ensure correct numerical behavior, and provides utilities for converting entire models between precisions without retraining.

Flax provides a collection of reference training loop implementations in examples/ covering common architectures (ResNet, Transformer, LSTM) and tasks (image classification, machine translation, language modeling). These examples demonstrate best practices for integrating Flax modules with Optax optimizers, Orbax checkpointing, and distributed training, serving as templates that users can fork and modify rather than framework features. The examples are intentionally simple and modular, encouraging users to customize training logic directly rather than relying on framework abstractions.