JAX for ML Systems

JAX is a Python library for numerical computing that turns ordinary array functions into compiled, differentiated, vectorized and sharded programs. This plan has 35 coding problems in nine sections, starting with how JAX arrays and random keys work and ending with the decoding loops and sharding patterns used to train and serve large models.

You should know Python and NumPy, and how gradient descent trains a model. Knowing how attention works in a Transformer helps in the later sections. The NumPy study plan and the attention mechanism are good preparation.

1. Arrays the JAX way: functional array updates, splitting random keys, masked reductions and padding to fixed bucket lengths.
2. Automatic differentiation: value and grad, Hessian-vector products, per-example gradient clipping and a straight-through estimator with a custom VJP.
3. Compiling with jit: static arguments, traced learning rate schedules, loops with scan and power iteration with a fixed number of steps.
4. Vectorizing with vmap: batched attention, model ensembles with in_axes and pairwise distances with nested vmap.
5. Pytrees and optimizers: SGD with momentum, global norm clipping, AdamW and gradient accumulation over nested parameter trees.
6. Memory and numerics: mixed precision with loss scaling, rematerialization, int8 per-channel quantization and blockwise attention with online softmax.
7. Inference building blocks: rotary embeddings, KV cache updates, top-k and top-p sampling and a grouped-query decode step.
8. Decoding loops: greedy decoding with scan, paged KV cache gathers, beam search and speculative decoding acceptance.
9. Sharding and parallelism: device meshes and named sharding, data parallelism with psum, tensor-parallel MLPs and a ring all-reduce.

Apply these patterns to model serving with the Inference Engineering study plan, or write the GPU kernels underneath them with the Triton study plan.