TPU Deep Dive
Course
Review
Slides
Exams
Sources
Theme: system
Sources
SB9
How to Profile TPU Programs
The TPU software stack
SB10
Programming TPUs in JAX
Auto, explicit, manual
J-shard
Sharding guide
J-jaxpr
jaxpr language
J-aot
Ahead-of-time lowering
J-log
JAX changelog
Shardy default 0.7.0; jax.shard_map 0.8.0
X-arch
XLA architecture
X-passes
HLO passes
X-dumps
HLO dumps
X-flags
Flags guidance
X-shardy-jax
Using Shardy from JAX
P-GSPMD
Xu et al., "GSPMD: General and Scalable Parallelization for ML Computation Graphs"
§3–4
N-graphs
CUDA Programming Guide §4.2, CUDA Graphs
§4.2
T-how
TorchTPU docs: How TorchTPU works
T-exec
TorchTPU docs: Execution model
T-dist
TorchTPU docs: Distributed support
T-kern
TorchTPU docs: Custom kernels
T-pt
PyTorch blog, "Introducing TorchTPU"
T-blog1
Google Developers Blog, "TorchTPU: Running PyTorch Natively on TPUs at Google Scale"
T-oss
Google Cloud Blog, "TorchTPU is now open source"
D-pytorch
PyTorch on TPU
T-xla
PyTorch/XLA repo
R-torchax
torchax (moved here 2025-10-06)
A-torch
PyTorch on Neuron / TorchNeuron
A-nki
NKI programming model
Chapter 9. Day 5, first block, 55 minutes
From Python to the MXU
Read
Quiz
12
Flashcards
21
Teach-back
6
Question 1 of 12
Recall
In what order does a jax.jit function pass through these representations on its way to a TPU?
1
jaxpr → StableHLO → HLO (XLA passes) → LLO → machine code
2
StableHLO → jaxpr → HLO → PTX
3
HLO → StableHLO → jaxpr → LLO
4
jaxpr → TorchInductor → LLO
Check answer
Enter
Press 1–4 to choose.