the museum
pytorch wing

serves the pytorch path · six errors reproduced on the course machine (torch 2.2.2, CPU), three captured on a Colab TPU v6e-1 during LAB·P4 and labelled where they sit · fixes run and proven · all wings

The tape objects, in its own words.

Nine failures. The first six come from the eager world, most of them autograd defending its tape: leaves mutated mid-graph, saved tensors rewritten, a graph walked twice. The last three come from the TPU bridges and could only be met on real hardware: the two bridges fighting over torch's dispatch, an install that left jax behind, and a random tensor whose key belonged to another client. Every snippet ran and failed; every error is verbatim (paths shortened); every fix ran and passed.

exhibit
01/09

Updating a leaf that autograd is watching

A leaf with requires_grad is a parameter in autograd's eyes, and mutating it mid-graph would corrupt the tape. Updates belong outside recording, under torch.no_grad(); that context is exactly how optimizer.step applies yours.

the failing code

import torch

x = torch.ones(3, requires_grad=True)
x += 1  # in-place on the leaf the optimizer would own

the error, verbatimRuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

the fix

import torch

x = torch.ones(3, requires_grad=True)
with torch.no_grad():
    x += 1  # what optimizer.step does under the hood

exhibit
02/09

Mutating a tensor the backward pass still needs

exp saves its output for the backward pass, and the in-place add rewrote that saved tensor. The version counter in the error is autograd noticing. Out-of-place ops leave the saved value alone; reserve in-place mutation for tensors no gradient will revisit.

the failing code

import torch

x = torch.ones(3, requires_grad=True)
y = x.exp()
y += 1  # exp's backward needs y; this rewrites it
y.sum().backward()

the error, verbatimRuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [3]], which is output 0 of ExpBackward0, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

the fix

import torch

x = torch.ones(3, requires_grad=True)
y = x.exp()
y = y + 1  # a new tensor; the saved output survives
y.sum().backward()

exhibit
03/09

Backward through a graph that was already freed

The tape is consumed by default: backward frees each node's saved tensors as it passes. A second backward finds nothing to walk. retain_graph=True keeps the graph alive at the cost of its memory; recomputing the forward is usually the better trade.

the failing code

import torch

x = torch.ones(3, requires_grad=True)
y = (x * x).sum()
y.backward()
y.backward()  # the tape was freed on the first pass

the error, verbatimRuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

the fix

import torch

x = torch.ones(3, requires_grad=True)
y = (x * x).sum()
y.backward(retain_graph=True)  # or recompute y; retaining costs memory
y.backward()

exhibit
04/09

backward on a tensor with more than one element

Reverse mode pulls one cotangent backward, so backward() with no argument only means something for a scalar. Reduce to a loss, or pass the cotangent explicitly. The JAX wing holds the same lesson wearing jax.grad's error message.

the failing code

import torch

x = torch.ones(3, requires_grad=True)
(x * 2).backward()

the error, verbatimRuntimeError: grad can be implicitly created only for scalar outputs

the fix

import torch

x = torch.ones(3, requires_grad=True)
(x * 2).sum().backward()  # reduce to a scalar, or pass a gradient argument

exhibit
05/09

A batch that does not fit the layer

nn.Linear(8, 4) owns an 8 by 4 weight, and the error is the matmul refusing 32x16 against it. Read the two shapes in the message: the second-to-last axis of yours must equal the layer's in_features, and one of the two declarations is wrong.

the failing code

import torch
from torch import nn

layer = nn.Linear(8, 4)
layer(torch.randn(32, 16))  # features axis is 16; the layer expects 8

the error, verbatimRuntimeError: mat1 and mat2 shapes cannot be multiplied (32x16 and 8x4)

the fix

import torch
from torch import nn

layer = nn.Linear(16, 4)  # declare the features the data actually has
layer(torch.randn(32, 16))

exhibit
06/09

Class targets handed over as floats

cross_entropy reads integer targets as class indices and float targets as probability rows, and a float vector of the wrong shape satisfies neither. Class labels travel as int64; the fix is the dtype, not the loss.

the failing code

import torch
from torch import nn

logits = torch.randn(4, 10)
targets = torch.tensor([1.0, 3.0, 2.0, 7.0])
nn.functional.cross_entropy(logits, targets)

the error, verbatimRuntimeError: expected scalar type Long but found Float

the fix

import torch
from torch import nn

logits = torch.randn(4, 10)
targets = torch.tensor([1, 3, 2, 7])  # class indices are int64
nn.functional.cross_entropy(logits, targets)

exhibit
07/09

Both bridges in one runtime

captured on a Colab TPU v6e-1 during LAB·P4, 2026-07-27; the fix ran green on the same hardware

torchax.enable_globally() is literal: it installs torch modes for the whole process, and torch_xla claims the same dispatch. With torchax still on, the tensors torch_xla hands to autograd carry no graph, so backward has nothing to walk. The chapter 10 sentence about two bridges is a statement about ownership, and this error is what sharing costs.

the failing code

import torchax
torchax.enable_globally()   # torch modes, installed process-wide

import torch_xla
model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1)).to('xla')
opt = torch.optim.Adam(model.parameters(), lr=1e-2)

with torch_xla.step():
    loss = nn.functional.mse_loss(model(x), y)
    opt.zero_grad()
    loss.backward()

the error, verbatimRuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

the fix

# one bridge per runtime. Take the modes down before switching:
import torchax
torchax.disable_globally()

# or, the version that always works: a fresh runtime holding only torch_xla
import torch_xla
with torch_xla.step():
    loss = nn.functional.mse_loss(model(x), y)
    opt.zero_grad()
    loss.backward()   # the graph is there again

exhibit
08/09

Installing torchax without its extra

captured on a Colab TPU v6e-1 during LAB·P4, 2026-07-27; the fix ran green on the same hardware

Bare torchax declares no jax dependency at all; the requirement lives in its extras, where the tpu extra asks for jax[tpu] and jax >= 0.6.2. Install it plain and pip leaves whatever jax the machine already had, which on a TPU VM that another install has moved can be too old to carry the internals torchax imports.

the failing code

!pip install -q torchax

import torchax   # whatever jax is already here is the jax you get

the error, verbatimImportError: cannot import name 'mutable_array' from 'jax.experimental'

the fix

!pip install -q 'torchax[tpu]'   # the extra pulls jax[tpu] and holds jax >= 0.6.2

import torchax
from jax.experimental import mutable_array   # make the version check explicit

exhibit
09/09

A random tensor whose key belongs to another client

captured on a Colab TPU v6e-1 during LAB·P4, 2026-07-27; the fix ran green on the same hardware

Creating straight onto the device routes through torchax's RNG, and its jax key belongs to whichever client was live when the key was made. On a machine carrying more than one jax or libtpu, the key and the array end up on different clients and the runtime says so. Building with plain torch and moving with .to('jax') never asks the question.

the failing code

inputs = torch.randn(4, 8, device='jax')   # torchax's own RNG makes this one

the error, verbatimXlaRuntimeError: INVALID_ARGUMENT: Unable to replace a PyArray with a PyArray from a different client.

the fix

inputs = torch.randn(4, 8).to('jax')   # build with plain torch, then move