Flax

Defines neural networks using functional programming patterns where module logic and state are strictly separated through the Scope system (flax/core/scope.py). Modules inherit from flax.linen.Module and implement __call__ methods that operate on immutable pytree state, enabling seamless composition with JAX transformations (jit, vmap, grad, pmap). State initialization happens explicitly via init() and inference via apply(), preventing hidden state mutations that cause JAX tracing errors.

Provides Python-native object-oriented module definitions (flax.nnx.Module) where parameters, buffers, and state are stored as instance attributes with automatic graph state management through GraphDef/State splitting (flax/nnx/graph.py). Modules use standard Python semantics (no explicit init/apply) while internally decomposing into a static computation graph (GraphDef) and mutable state (State) that can be independently transformed. This bridges imperative programming familiarity with JAX's functional requirements.

Implements a variable collection system (flax/core/scope.py, flax/linen/module.py) that tracks different types of model state (params, cache, batch_stats, dropout_rng) separately through the Scope abstraction. Variables are collected into named collections that can be selectively updated or frozen during training. For example, batch normalization statistics are tracked in 'batch_stats' collection and updated separately from parameters. This enables fine-grained control over which state is updated during training vs. inference.

Integrates JAX's automatic differentiation (jax.grad, jax.value_and_grad) with Flax's state management to enable efficient gradient computation through jit-compiled training steps. Gradients are computed with respect to parameters while preserving other state (batch_stats, cache) through mutable variable collections. Integration with Optax optimizers enables atomic parameter updates with momentum, adaptive learning rates, and gradient clipping. Training steps are typically jit-compiled for performance, with gradients computed and applied in a single compiled function.

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.

Wraps JAX transformations (jit, vmap, grad, pmap, scan) with Flax-aware variants (flax/core/lift.py, flax/linen/transforms.py) that automatically handle variable collection and state threading through transformation boundaries. For example, nn.vmap maps over batch dimensions while preserving parameter sharing across mapped instances, and nn.scan unrolls recurrent operations while managing hidden state across timesteps. These lifted transforms eliminate manual state threading boilerplate that would otherwise be required.

Implements single-program-multiple-data (SPMD) parallelism through JAX's pmap and sharding APIs, with Flax-specific utilities for annotating model parameters and activations with sharding constraints (flax/linen/transforms.py, distributed training utilities). Developers specify logical axis names (e.g., 'batch', 'heads', 'vocab') and Flax automatically generates sharding directives that map to physical device mesh topology. This abstracts away low-level pmap complexity while enabling multi-host, multi-device training without code changes.

Provides flax.training.train_state.TrainState, a pytree container that bundles model parameters, optimizer state, and training metadata (step count, learning rate schedule) into a single immutable structure. TrainState integrates with Optax optimizers to provide a standard training loop pattern: state = train_step(state, batch) where train_step applies gradients and updates optimizer state atomically. This eliminates manual state threading and provides a consistent interface across different optimization algorithms.

Integrates with Orbax (Google's checkpointing library) to provide flax.training.checkpoints utilities for saving/loading model parameters, optimizer state, and training metadata to disk or cloud storage. Supports multiple serialization formats (msgpack, pickle, safetensors) and enables asynchronous checkpointing that doesn't block training. Flax checkpoints are pytrees, enabling efficient incremental saves and restoration of distributed training state across device topologies.

Provides a comprehensive library of neural network layers (Dense, Conv, Attention, LayerNorm, Dropout, etc.) implemented in both Linen (flax/linen/nn/) and NNX (flax/nnx/nn/) APIs, with JAX-specific optimizations like fused operations and efficient attention implementations. Layers are composable building blocks that handle parameter initialization, shape inference, and numerical stability automatically. Attention layers use efficient kernels (e.g., flash attention patterns) and support multi-head, multi-query, and grouped query variants.

Provides utilities (flax.linen.summary, module introspection APIs) to inspect model structure, count parameters, estimate memory usage, and generate model summaries without running forward passes. Introspection works by analyzing the module graph structure and parameter pytrees, enabling developers to understand model complexity before training. Summary output shows layer-by-layer parameter counts, shapes, and computational costs.

Provides reference implementations of complete training loops (flax/examples/) for common tasks (image classification, sequence-to-sequence, language modeling) that demonstrate best practices for data loading, gradient computation, metric tracking, and checkpoint management. These examples are designed to be forked and modified rather than used as black-box APIs, enabling researchers to customize training logic without fighting framework abstractions. Examples cover single-device, multi-device, and distributed training patterns.

Implements a two-phase initialization protocol where modules are first initialized with dummy inputs to infer parameter shapes, then applied with actual data. Initialization is handled through flax.linen.Module.init() which returns a pytree of parameters, enabling shape inference without manual specification. This approach ensures type safety and prevents shape mismatches at runtime. Initialization can be customized through kernel_init and bias_init functions that specify parameter distributions (e.g., normal, uniform, orthogonal).