tpu

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

Trace a two-by-two output-stationary array cycle by cycle. Derive its fill and drain cost without turning a toy schedule into a TPU benchmark.

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.

Validate the exact model, accelerator generation and serving features. Separate packed storage from kernel execution, derive the real memory fallback, and compare qualified throughput and total migration cost.

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.