[cokernel_hacker]

This python actually builds a graph under the hood which then gets JIT compiled for CPU/GPU/TPU.

[senseiV]

Does a TPU have XLA-graph for GPUs Cuda-graphs? Not sure on TPU theory

[dekhn]

TPUs lower XLA into their own instruction set. https://pytorch.org/xla/release/2.3/index.html