flax

Flax provides a module system built on JAX's functional programming paradigm, allowing developers to define neural networks as composable classes that separate model definition from parameter state. Modules use a two-phase initialization pattern: first defining architecture through class inheritance, then materializing parameters through explicit initialization calls that return immutable pytrees. This design enables automatic differentiation through JAX's jit, grad, and vmap transformations without stateful mutation.

Flax implements lazy parameter initialization where module shapes are inferred at first forward pass rather than requiring explicit shape specification upfront. The framework traces through the model with dummy input arrays to discover parameter dimensions, then materializes the full parameter tree in a single initialization call. This eliminates manual shape calculation and supports dynamic architectures where layer sizes depend on input dimensions.

Flax provides utilities for gradient checkpointing (also called activation checkpointing) that trade computation for memory by recomputing activations during backpropagation instead of storing them. This enables training larger models on memory-constrained devices. The framework also supports gradient accumulation where gradients are computed over multiple batches before updating parameters, enabling larger effective batch sizes without proportional memory increases.

Flax integrates with JAX's mixed precision capabilities to enable training with lower-precision computations (float16, bfloat16) while maintaining numerical stability through loss scaling. Loss scaling prevents gradient underflow by multiplying losses before backpropagation, then unscaling gradients before parameter updates. The framework provides utilities for automatic loss scaling that dynamically adjusts the scale factor based on gradient overflow detection.

Flax provides patterns and utilities for distributed training across multiple devices (GPUs, TPUs) using JAX's pmap (parallel map) and pjit (parallel jit) primitives. These enable data parallelism (splitting batches across devices) and model parallelism (splitting parameters across devices) without requiring manual communication code. The framework includes examples and utilities for common distributed patterns (data parallelism, pipeline parallelism) that work seamlessly with Flax's functional training loops.

Flax provides training utilities that wrap JAX's grad and jit transformations into reusable patterns, handling parameter updates, loss computation, and metric aggregation without requiring manual gradient tape management. The framework uses a TrainState abstraction that bundles parameters, optimizer state, and step count into a single pytree, enabling clean functional updates through optimizer.apply_gradients() calls. Metrics are computed as pure functions and aggregated across batches through pytree operations.

Flax provides production-ready implementations of multi-head attention, transformer blocks, and positional encodings optimized for numerical stability and JAX compatibility. Attention uses log-space softmax computation to prevent overflow, supports arbitrary query/key/value projections, and integrates with JAX's vmap for efficient batch processing. Transformer blocks compose attention, feed-forward networks, and layer normalization with configurable residual connections and dropout patterns.

Flax provides checkpoint utilities that serialize model parameters and optimizer state as pytrees to disk, supporting multiple formats (pickle, msgpack, SafeTensors) with automatic compression and versioning. The framework includes utilities for partial checkpointing (saving only parameters, only optimizer state, or both), resuming training from checkpoints with state reconstruction, and loading pre-trained weights into models with different architectures through flexible key matching.

Flax implements batch and layer normalization layers that correctly handle training vs. inference modes through explicit state management. During training, batch statistics are computed and accumulated into a running statistics buffer; during inference, accumulated statistics are used. The framework uses a mutable state pattern where normalization layers return both outputs and updated statistics, which are merged back into the model state after each forward pass.

Flax implements dropout and other stochastic layers using JAX's functional random number generation, where RNG keys are threaded through the model as explicit parameters rather than using global random state. During training, dropout masks are generated from the RNG key; during inference, dropout is disabled. The framework provides utilities for splitting RNG keys across batches and layers, ensuring reproducibility and correct behavior with jit compilation.

Flax provides embedding layers that map discrete token indices to dense vectors, with support for weight sharing between input embeddings and output projection layers (useful in language models). Embeddings are stored as parameter matrices and support arbitrary vocabulary sizes and embedding dimensions. The framework includes utilities for initializing embeddings with specific distributions and for sharing weights across multiple embedding instances.

Flax implements RNN, LSTM, and GRU cells as stateless modules that process one timestep at a time, returning both output and updated hidden state. Sequences are processed by manually iterating over timesteps and threading the hidden state through each step, or using scan operations for efficient batched processing. This design maintains functional purity while supporting arbitrary sequence lengths and enables efficient compilation through JAX's scan primitive.

Flax provides convolutional layers (Conv1D, Conv2D, Conv3D) that apply learned filters to input arrays with configurable kernel sizes, strides, padding modes, and dilation. Convolutions are implemented as efficient JAX operations that leverage underlying BLAS libraries and support arbitrary input/output channel dimensions. The framework includes utilities for grouped convolutions and depthwise separable convolutions for parameter efficiency.