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
21
left in this session,
21
of them new.
Reviewed
0
Question
Compilation stages from JAX to TPU?
Show answer
Space