Flax NNX MLP
A small multilayer perceptron with nested NNX module scopes.
JAXViz
Model examples
Explore JAX forward passes alongside the code that produced them. Distributed examples include both the logical global program and the partitioned per-device program.
Foundations
Flax Linen MLP
The same basic architecture expressed with Flax Linen.
Equinox MLP
A compact Equinox model traced as a plain JAX callable.
Distributed MLPs
Tensor-parallel MLP (plain JAX)
The smallest model that needs communication: one matrix split by columns, the next by rows.
Tensor-parallel MLP (Flax NNX)
Column- and row-parallel linear layers across one model axis.
2D-parallel MLP
Data parallelism and tensor parallelism combined on a 2 × 4 mesh.
Attention
Head-parallel self-attention
Batch sharding across data devices and complete attention heads across model devices.
Transformer encoder
Four pre-norm blocks for selectively expanding attention and feed-forward subgraphs.