Flax is a high-level neural-network library for JAX. Its Linen API combines familiar module-based modeling with JAX’s functional transformations, so you can assemble reusable dense, convolutional, LSTM, attention, and activation-function blocks without implementing everything from low-level primitives. section 11.0 section 11.4
A practical Flax workflow separates model definition from state handling:
- Define a network with Linen Modules, using an object-oriented style inside a functional system. section 11.4
- Initialize parameters and other model variables with the instantiated Module’s init() method. section 11.4
- Run a forward pass with apply(), passing variables and input data explicitly. section 11.4
- Represent simple training state with TrainState and pair training with Optax gradient transformations. section 11.0 section 11.4
A model pass
Create a Linen Module, call init() to obtain its variables, then call apply() with those variables and an input batch. The same explicit handling extends to internal state such as BatchNorm state, while TrainState can collect parameters, optimizer information, and metrics for a training loop. section 11.4 section 11.0
