Why does my torch_xla notebook stop working weeks after it ran fine?

Almost always the version pair drifted. torch and torch_xla ship as a matched set and the two numbers have to agree: an install line that pins one and floats the other will, sooner or later, pick up a torch that the installed torch_xla was never built against, and the failure shows up as an import error or a missing symbol rather than as anything mentioning versions. Pin both, in one line, to the same number. The labs on this track use torch==2.9.0 with torch_xla[tpu]==2.9.0, and 2.9.0 is what setup.py stamps at the commit they were read against. The second cause is the accelerator wheel underneath: the [tpu] extra pulls a libtpu pinned to that release, and anything else in the environment that also wants libtpu (jax[tpu], for instance) will move it. A notebook that installs one stack per runtime and never mixes two survives; a notebook that installs on top of a runtime that already has an accelerator stack is rolling dice.

What does PJRT_DEVICE do, and what happens if I leave it unset?

PJRT_DEVICE names which runtime backend torch_xla should build: CPU, TPU, NEURON, XPU, or a dynamically registered plugin. It is read on the C++ side the first time anything needs a device, and the client that comes back is built once per process and never rebuilt. Leaving it unset does not fail, which is the part that surprises people: the last line of torch_xla/__init__.py calls a default-selection helper that picks TPU when a libtpu and real chips are both present, NEURON when the Neuron plugin is installed, and otherwise logs "Defaulting to PJRT_DEVICE=CPU" and sets CPU. So a notebook on a Colab TPU runtime where the TPU was never actually attached will quietly train on CPU while you wonder why it is slow, and the only evidence is a warning line in the import output. Set it explicitly at the top of the notebook so the choice lives in the file. The hard error, "$PJRT_DEVICE is not set.", only fires when the variable is present but empty, or when the default selection has been switched off.

Why is the first training step so much slower than the rest?

Because the first step is where the compile happens. On the xla device your ops do not run when you write them; they record. torch_xla.sync() cuts the recorded graph, hashes it, looks the hash up in a compilation cache, and only on a miss does it hand the program to the compiler. That compile is the tail on step one. Step two, with identical shapes, produces an identical hash, hits the cache, and skips straight to execution. You can see the split rather than infer it: the metrics report keeps CompileTime and ExecuteTime as separate timers, and the counters CachedCompile and UncachedCompile say which way each lookup went. If UncachedCompile keeps climbing after the first few steps, something in your loop is changing shapes, and every changed shape is a new hash and a new compile.

Can I use torchax and torch_xla in the same notebook?

No, and the reason it is worth knowing is that it fails twice, in two different places. At run time both libraries claim PyTorch’s dispatcher (torchax installs global torch modes, torch_xla registers a lazy backend at import), and the loser’s tensors come out with no autograd graph attached, so loss.backward() raises "element 0 of tensors does not require grad and does not have a grad_fn", which reads like a bug in your model and is not. At install time they fight over libtpu: torch_xla[tpu] pins one version, jax[tpu] (which torchax needs) brings another, and pip resolves that by moving jax, which is how a torchax import that worked ten minutes ago starts failing. Restarting the session does not help, because a restart keeps every installed file on disk. The fix is to delete the runtime, not restart it, and start the second bridge in a clean one.

How do I see the graph torch_xla built without running it?

Two calls on the compiled extension print the pending graph in place: _XLAC._get_xla_tensors_text([t]) walks the IR nodes behind a tensor and prints them, and _XLAC._get_xla_tensors_hlo([t]) lowers the same thing to readable HLO. Neither executes anything and neither consumes the pending graph, so the tensor is still un-materialized afterwards and the next sync still compiles and runs it. For a record you can keep, set XLA_SAVE_TENSORS_FILE to a path and torch_xla appends every graph it cuts to that file, in the format XLA_SAVE_TENSORS_FMT selects (text, hlo, or dot); add XLA_IR_DEBUG=1 and XLA_HLO_DEBUG=1 to get the Python frames that built each node carried into the dump. One trap: the path you set is not the path that appears. torch_xla appends a dot and the device ordinal, so XLA_SAVE_TENSORS_FILE=/tmp/trace.ir writes /tmp/trace.ir.0.

Writing torch_xla Notebooks That Survive Colab

A Colab notebook that trains a model on a TPU is easy to get running once and surprisingly hard to get running twice. This is the craft layer under torch_xla: why torch and torch_xla are one pinned pair and never two independent versions; what PJRT_DEVICE actually does when you leave it unset (import picks a default and tells you in a warning line most people scroll past); why torchax and torch_xla cannot share a runtime, at run time over PyTorch’s dispatcher and again at install time over libtpu; why the first step is slow and the second one is not; and the three instruments that let you prove any of it: the metrics report, the IR and HLO dumps written to a file whose name is not the name you gave it, and the two calls that print a graph without running it. Closes with a cell order that works and a paste-back discipline that keeps a reference run honest.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Writing torch_xla Notebooks That Survive Colab — built to be understood, not skimmed.

Questions

Why does my torch_xla notebook stop working weeks after it ran fine?: Almost always the version pair drifted. torch and torch_xla ship as a matched set and the two numbers have to agree: an install line that pins one and floats the other will, sooner or later, pick up a torch that the installed torch_xla was never built against, and the failure shows up as an import error or a missing symbol rather than as anything mentioning versions. Pin both, in one line, to the same number. The labs on this track use torch==2.9.0 with torch_xla[tpu]==2.9.0, and 2.9.0 is what setup.py stamps at the commit they were read against. The second cause is the accelerator wheel underneath: the [tpu] extra pulls a libtpu pinned to that release, and anything else in the environment that also wants libtpu (jax[tpu], for instance) will move it. A notebook that installs one stack per runtime and never mixes two survives; a notebook that installs on top of a runtime that already has an accelerator stack is rolling dice.
What does PJRT_DEVICE do, and what happens if I leave it unset?: PJRT_DEVICE names which runtime backend torch_xla should build: CPU, TPU, NEURON, XPU, or a dynamically registered plugin. It is read on the C++ side the first time anything needs a device, and the client that comes back is built once per process and never rebuilt. Leaving it unset does not fail, which is the part that surprises people: the last line of torch_xla/__init__.py calls a default-selection helper that picks TPU when a libtpu and real chips are both present, NEURON when the Neuron plugin is installed, and otherwise logs "Defaulting to PJRT_DEVICE=CPU" and sets CPU. So a notebook on a Colab TPU runtime where the TPU was never actually attached will quietly train on CPU while you wonder why it is slow, and the only evidence is a warning line in the import output. Set it explicitly at the top of the notebook so the choice lives in the file. The hard error, "$PJRT_DEVICE is not set.", only fires when the variable is present but empty, or when the default selection has been switched off.
Why is the first training step so much slower than the rest?: Because the first step is where the compile happens. On the xla device your ops do not run when you write them; they record. torch_xla.sync() cuts the recorded graph, hashes it, looks the hash up in a compilation cache, and only on a miss does it hand the program to the compiler. That compile is the tail on step one. Step two, with identical shapes, produces an identical hash, hits the cache, and skips straight to execution. You can see the split rather than infer it: the metrics report keeps CompileTime and ExecuteTime as separate timers, and the counters CachedCompile and UncachedCompile say which way each lookup went. If UncachedCompile keeps climbing after the first few steps, something in your loop is changing shapes, and every changed shape is a new hash and a new compile.
Can I use torchax and torch_xla in the same notebook?: No, and the reason it is worth knowing is that it fails twice, in two different places. At run time both libraries claim PyTorch’s dispatcher (torchax installs global torch modes, torch_xla registers a lazy backend at import), and the loser’s tensors come out with no autograd graph attached, so loss.backward() raises "element 0 of tensors does not require grad and does not have a grad_fn", which reads like a bug in your model and is not. At install time they fight over libtpu: torch_xla[tpu] pins one version, jax[tpu] (which torchax needs) brings another, and pip resolves that by moving jax, which is how a torchax import that worked ten minutes ago starts failing. Restarting the session does not help, because a restart keeps every installed file on disk. The fix is to delete the runtime, not restart it, and start the second bridge in a clean one.
How do I see the graph torch_xla built without running it?: Two calls on the compiled extension print the pending graph in place: _XLAC._get_xla_tensors_text([t]) walks the IR nodes behind a tensor and prints them, and _XLAC._get_xla_tensors_hlo([t]) lowers the same thing to readable HLO. Neither executes anything and neither consumes the pending graph, so the tensor is still un-materialized afterwards and the next sync still compiles and runs it. For a record you can keep, set XLA_SAVE_TENSORS_FILE to a path and torch_xla appends every graph it cuts to that file, in the format XLA_SAVE_TENSORS_FMT selects (text, hlo, or dot); add XLA_IR_DEBUG=1 and XLA_HLO_DEBUG=1 to get the Python frames that built each node carried into the dump. One trap: the path you set is not the path that appears. torch_xla appends a dot and the device ordinal, so XLA_SAVE_TENSORS_FILE=/tmp/trace.ir writes /tmp/trace.ir.0.

Writing torch_xla Notebooks That Survive Colab

Questions

Related explainers