jax

Carry compute and memory reasoning to TPU v6e, correct all-gather ownership arithmetic, and learn the compilation, sharding and custom-kernel workflow.

A global-array program lets the compiler derive communication from sharding choices. Work through a small matrix example, then separate that programming model from TPU topology, compilation and Pathways recovery.

Three routes with different costs: keep PyTorch through XLA, bet on the native backend, or adopt the JAX reference trainer. Choose from the model's custom code and the timeline, then prove parity before scaling.

Same model, same engine name, different execution model. Trace what moves into compilation, how shape buckets trade compile time for padding, and how per-device memory decides the slice.