JAX Training Recipe

JAX is not PyTorch with different syntax. Functions are pure, state is explicit, and transformations (jit, vmap, grad, pmap) compose only when the function boundaries are clean.

For concrete config/checklist details, read references/training-recipe-details.md.

Intake

Specify:

• task, dataset, loss, and metric;
• target device: CPU, single GPU, multi-GPU, TPU;
• model framework: Flax/Linen, Equinox, or raw JAX;
• randomness needs: dropout, sampling, augmentation;
• state needs: batch norm, optimizer state, EMA, checkpointing.

Workflow