2. Tensor operations and autograd mechanics

Every PyTorch model, from a two-layer perceptron to a transformer, is ultimately a composition of tensor operations. What distinguishes PyTorch from a plain numerical library is that these operations are recorded on a dynamic computation graph, and the autograd engine can differentiate through them automatically. Understanding both layers — the tensor algebra and the gradient machinery — is what allows you to write custom layers, debug silent shape bugs, and reason about why a loss is not decreasing.

Tensors as typed, device-aware arrays

A tensor is a multidimensional array with three properties that matter for model development: shape, dtype, and device. Shape determines which operations are legal; dtype determines numerical precision and memory cost; device determines whether an operation runs on CPU or GPU. A frequent source of runtime errors is a mismatch in any of these three.

import torch

x = torch.randn(32, 3, 224, 224)      # NCHW image batch, float32, CPU
w = torch.randn(64, 3, 7, 7)          # conv weight: out_ch, in_ch, kH, kW
print(x.shape, x.dtype, x.device)
# torch.Size([32, 3, 224, 224]) torch.float32 cpu

y = torch.nn.functional.conv2d(x, w, stride=2, padding=3)
print(y.shape)                        # torch.Size([32, 64, 112, 112])

Broadcasting follows NumPy semantics: dimensions are aligned from the right, and a dimension of size 1 is stretched to match. This is convenient but dangerous. Adding a tensor of shape (32, 1) to one of shape (32,) produces (32, 32), not (32,). When a loss suddenly has an unexpected magnitude, a silent broadcast is a common culprit. Use torch.broadcast_shapes(a.shape, b.shape) or simply assert shapes before arithmetic in critical code paths.

Views, reshape, and memory layout

Operations such as view, transpose, permute, and squeeze do not copy data; they return a new tensor object that shares the same storage with different strides. reshape returns a view when possible and copies otherwise. This distinction matters because in-place writes to a view mutate the original tensor, and because a non-contiguous tensor cannot always be passed to view.

a = torch.arange(12).reshape(3, 4)
b = a.t()                 # shape (4, 3), non-contiguous
# b.view(12)              # RuntimeError: view size is not compatible
c = b.reshape(12)         # OK: copies into contiguous memory
d = b.contiguous().view(12)  # explicit copy, then view

In practice, prefer reshape unless you specifically need a view, and call .contiguous() before passing a transposed tensor to a kernel that requires contiguous input.

How autograd records the graph

When requires_grad=True is set on a tensor, every operation involving it creates a node in a directed acyclic graph. Each node stores the forward output and a backward function that knows how to compute the gradient with respect to its inputs. Calling loss.backward() traverses this graph in reverse topological order, applying the chain rule.

x = torch.tensor([2.0, 3.0], requires_grad=True)
w = torch.tensor([1.5, -0.5], requires_grad=True)
b = torch.tensor(0.1, requires_grad=True)

z = (x * w).sum() + b        # forward: z = 2*1.5 + 3*(-0.5) + 0.1 = 1.6
z.backward()

print(x.grad)   # tensor([ 1.5000, -0.5000])  -> dz/dx = w
print(w.grad)   # tensor([2., 3.])            -> dz/dw = x
print(b.grad)   # tensor(1.)                  -> dz/db = 1

Three rules govern this process. First, gradients accumulate by default: calling backward() twice without zeroing adds the second gradient to the first. Second, the graph is freed after backward() unless retain_graph=True is passed, which is why you cannot call backward() twice on the same loss without it. Third, only leaf tensors with requires_grad=True populate .grad; intermediate results do not, unless you call retain_grad() on them.

The training step as a closed loop

The canonical training iteration ties tensors, autograd, and the optimizer together. The order of operations is not arbitrary: zeroing must precede the backward pass, and the optimizer step must follow it.

model = torch.nn.Linear(4, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for step in range(3):
    xb = torch.randn(8, 4)
    yb = torch.randn(8, 1)

    pred = model(xb)
    loss = torch.nn.functional.mse_loss(pred, yb)

    optimizer.zero_grad()   # 1. clear old gradients
    loss.backward()         # 2. compute new gradients
    optimizer.step()        # 3. update parameters

    print(step, float(loss))

To verify that gradients are correct, compare them against a finite-difference approximation. This is the standard sanity check when implementing a custom autograd.Function or a new loss.

from torch.autograd import gradcheck

f = lambda inp: (inp ** 3).sum()
inp = torch.randn(5, dtype=torch.double, requires_grad=True)
print(gradcheck(f, (inp,), eps=1e-6, atol=1e-4))   # True if gradients match

Controlling gradient flow

Three mechanisms let you control which parts of the graph participate in differentiation. torch.no_grad() disables graph recording entirely and is used during evaluation and inference to save memory and time. detach() returns a tensor that shares data but is severed from the graph, useful for targets and for stopping gradients through a branch. torch.autograd.grad() computes gradients without populating .grad, which is convenient for meta-learning and second-order methods.

with torch.no_grad():
    eval_pred = model(xb)          # no graph built

target = yb.detach()               # target never needs gradients
branch = pred.detach()             # gradient stops here

For higher-order gradients, pass create_graph=True to backward() or use torch.autograd.grad with create_graph=True. This is required for gradient penalty terms in WGAN-GP and for MAML-style inner loops.

Diagnosing common failures

When a model fails to learn, the symptoms usually point to a specific autograd or tensor issue. The table below maps observable behavior to likely causes and the first check to run.

What to carry forward

Tensors give you the data representation; autograd gives you the derivative of any scalar with respect to any leaf. The training loop is the contract that binds them: zero, forward, backward, step. Once you can read a computation graph, predict gradient shapes, and recognize the difference between a view and a copy, you are ready to build custom modules and losses that the autograd engine can differentiate correctly.