From Python to the MXU

Tick each objective you can deliver out loud without notes. Then write the one-paragraph explanation you would give a room of GPU engineers. Saved attempts stay below, so you can compare them over the week.

• I can trace a function from Python through jaxpr, StableHLO, HLO and XLA's passes to LLO and TPU instructions, and dump each stage.
• I can say what I control on TPU — shardings, constraints, layouts, remat, donation, flags and Pallas — and what I do not.
• I can explain Shardy against GSPMD, compare JAX's three sharding modes (auto, explicit, manual shard_map), and say when to drop to each.
• I can compare XLA whole-program compilation with CUDA Graphs and torch.compile with Inductor.
• I can explain TorchTPU in detail — the PrivateUse1 backend, deferred StableHLO ops, its three eager modes, torch.compile without Inductor, the compile cache, tpu_dist and the SPMD handshake, and Pallas through jax_op.
• I can contrast TorchTPU with JAX, torchax and PyTorch/XLA, and the TPU stack with AWS's Neuron stack.

Prompts to answer out loud

1. A CUDA engineer says 'XLA is just CUDA Graphs.' Correct them in three sentences.
2. Draw how y = model(x) reaches the TPU under JAX jit, torchax, TorchTPU strict eager and TorchTPU torch.compile.
3. Why did TorchTPU need an 'SPMD handshake' that JAX never needed?
4. XLA inserted a giant AllGather you did not expect. What are your three options, from least to most effort?