PyTorch Tensors and Autograd for People Who Wrote Backprop by Hand

The NumPy examples in this series built forward computations and then propagated derivatives through them by hand. PyTorch can record supported tensor operations during the forward pass and apply their backward rules automatically. The chain rule is the same; the framework handles bookkeeping across paths and parameters. Understanding those paths still helps when implementing custom operations or diagnosing an unexpected gradient.

From NumPy arrays to PyTorch tensors

A torch.Tensor is a multidimensional array with a shape, dtype, and device. Much of its indexing, broadcasting, and matrix multiplication will be familiar from NumPy, although APIs and some defaults differ. The examples below run in order with PyTorch installed. They use CPU tensors until the final device example.

import torch
import numpy as np

# This supported CPU tensor and its NumPy array share storage.
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
b = torch.ones(2, 2)
print(a.shape, a.dtype, a.device)
print((a @ b).shape)
print(a.sum(dim=0, keepdim=True).shape)

arr = a.numpy()
arr[0, 0] = 99.0
print(a[0, 0])
# torch.Size([2, 2]) torch.float32 cpu
# torch.Size([2, 2])
# torch.Size([1, 2])
# tensor(99.)

Reductions use dim where NumPy commonly uses axis, and keepdim=True preserves a reduced axis. In the example, summing two rows along dimension 0 gives one row of two column sums. Python floating-point values normally create float32 PyTorch tensors; NumPy normally creates float64 arrays. Specify the dtype when precision matters. Numerical gradient checks are usually more reliable in double precision, but lower precision does not make every check meaningless.

The NumPy conversion above shares CPU storage, so changing arr changes a. This applies to the supported CPU tensor used here, which does not require gradients. torch.from_numpy(array) also shares supported NumPy storage; torch.tensor(array) copies it. To export a gradient-tracked tensor, use tensor.detach().cpu().numpy(); append .copy() if the array must be independent. Avoid modifying shared arrays while autograd may still need their values.

The chain-rule machinery behind every gradient here is derived in Derivatives and Computation Graphs for Neural Network Learning.

requires_grad builds the graph

For floating-point tensors in ordinary grad mode, requires_grad=True enables tracking through operations with backward rules. A leaf is a tensor created directly for differentiation, such as a model parameter, rather than a result of another tracked operation. Calling backward() on a real scalar result accumulates its derivatives in the participating leaves’ .grad fields. Comparisons and integer outputs do not acquire a differentiable path merely because their inputs require gradients.

x = torch.tensor(3.0, requires_grad=True)
y = torch.tensor(4.0, requires_grad=True)
z = x * y
z.retain_grad()
w = z + x ** 2
w.backward()
print(x.grad, y.grad)
print("leaf flags", x.is_leaf, z.is_leaf)
print("z has backward history", z.grad_fn is not None)
print("retained z gradient", z.grad)
# tensor(10.) tensor(3.)
# leaf flags True False
# z has backward history True
# retained z gradient tensor(1.)

Here \(x\) contributes through both \(xy\) and \(x^2\), giving \(4+6=10\); \(y\) contributes through the product, giving 3. z.grad_fn identifies its backward operation, while is_leaf distinguishes inputs from intermediate results. Intermediate tensors normally do not retain a .grad value. Calling z.retain_grad() before backward requests it; here \(\partial w/\partial z=1\).

The forward pass records the operations that actually execute, including the selected branches of Python control flow. It does not differentiate the decision to take one discrete branch rather than another. After backward, saved values needed by that backward calculation are normally released. Reusing the same result for another backward therefore often raises an error; computing a new forward result builds a new graph. retain_graph=True preserves saved data when reuse is needed, at a memory cost. It is not the usual solution for successive training steps.

Gradients accumulate — that is deliberate

Backward calls on fresh forward results add to existing leaf gradients. Clear gradients at the beginning of each intended accumulation window, often with optimizer.zero_grad(set_to_none=True). With one batch per update, that window contains one backward pass. Deliberate micro-batch accumulation keeps the parameters fixed until all contributions have been added. Forgetting to clear between unrelated updates changes the update rule; it does not guarantee either slow learning or decreasing loss.

p = torch.tensor(2.0, requires_grad=True)
(p ** 2).backward()
print(p.grad)
(p ** 2).backward()
print(p.grad)
p.grad = None
(p ** 2).backward()
print(p.grad)
# tensor(4.)
# tensor(8.)
# tensor(4.)

The repeated p ** 2 expressions above each create a fresh graph. Across micro-batches, weighting matters: if their losses are means over \(b_j\) examples and the total is \(B\), multiply each by \(b_j/B\) before backward to obtain the mean gradient for the window, assuming the loss separates across examples and the forward computation is equivalent. BatchNorm or other batch-dependent computations can make micro-batches differ from one large batch. Setting .grad=None releases the buffer; zeroing an existing buffer reuses it. Optimizers may skip a parameter whose gradient is None, so this distinction can matter for unused parameters.

When the output is a vector

A vector output does not specify one scalar objective. Supply an upstream gradient of the same shape, or reduce the vector to a scalar first. For \(v=x^2\), v.backward(weights) computes the gradient of \(\sum_i\text{weights}_i v_i\). With \(x=[2,3]\) and weights [1,2], the result is \([2\times2\times1,2\times3\times2]=[4,12]\). This is a vector–Jacobian product, not the full matrix of every output derivative.

vector_input = torch.tensor([2., 3.], requires_grad=True)
vector_output = vector_input ** 2
vector_output.backward(torch.tensor([1., 2.]))
print(vector_input.grad)
# tensor([ 4., 12.])

Turning the graph off

Autograd saves values required by the backward rules; it does not save every intermediate tensor indiscriminately. torch.no_grad() prevents backward-graph recording for operations in its scope. Ordinary prediction can use it to reduce autograd overhead. Evaluation that needs input gradients, such as a saliency calculation, must keep gradient recording enabled.

torch.manual_seed(1)
model = torch.nn.Linear(2, 2)
batch = torch.ones(3, 2)
model.eval()
tracked = model(batch)
scores = tracked.detach()
with torch.no_grad():
    preds = model(batch)
print("requires_grad", tracked.requires_grad, scores.requires_grad, preds.requires_grad)
print("same predictions", torch.allclose(tracked, preds))
# requires_grad True False False
# same predictions True

The tracked forward pass above has already built a graph before detach() is called. Detaching returns a tensor without that history; it does not avoid the original recording cost, and it shares storage with the source. no_grad avoids recording the forward in the first place. model.eval() is separate: it changes layers such as dropout and BatchNorm, but does not disable autograd. Ordinary inference commonly uses both evaluation mode and no_grad.

Logistic regression, ported

This logistic regression example uses rows as examples: \(X\) has shape \((m,n)\), \(w\) has shape \((n,)\), and the logits and labels each have shape \((m,)\). The target is generated from a noisy linear score. The loss averages binary cross-entropy over the \(m\) examples. Autograd supplies the same derivatives as the earlier NumPy implementation, while parameter updates remain explicit.

torch.manual_seed(0)
m, n = 200, 5
X = torch.randn(m, n)
true_w = torch.tensor([1.5, -2.0, 0.5, 0.0, 1.0])
y = (X @ true_w + 0.1 * torch.randn(m) > 0).float()
w = torch.zeros(n, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.5
for step in range(300):
    logits = X @ w + b
    loss = torch.nn.functional.binary_cross_entropy_with_logits(logits, y)
    w.grad = None
    b.grad = None
    loss.backward()
    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad
    if step % 100 == 0:
        print(step, round(loss.item(), 4))
with torch.no_grad():
    acc = (((X @ w + b) > 0).float() == y).float().mean()
print("training accuracy", round(acc.item(), 4))
# 0 0.6931
# 100 0.1622
# 200 0.1277
# training accuracy 0.975

loss.backward() computes \(\nabla_w L=X^\top(\operatorname{sigmoid}(Xw+b)-y)/m\) and the corresponding mean residual for the bias. The update runs under no_grad so that the in-place changes to gradient-tracked leaves are allowed and are not recorded as part of a new graph. Do not overwrite values between a forward pass and its backward if the derivative needs those values. The fused binary_cross_entropy_with_logits accepts raw logits and evaluates the loss stably without taking logs of rounded sigmoid probabilities.

The printed losses are measured before their corresponding updates; the accuracy is measured after all 300 updates. A logit above zero corresponds to a sigmoid probability above 0.5. This is training accuracy on the same synthetic examples used for fitting, not an estimate from held-out data. .item() extracts a Python number from a one-element tensor for reporting.

Checking derivatives numerically

When a function is built from PyTorch operations, torch.autograd.gradcheck compares its autograd derivatives with finite differences. For a custom autograd.Function, it instead checks the backward rule you supplied. The default tolerances are designed for double precision. The fixed inputs below are away from the kink of abs at zero; nondifferentiable points can make a numerical comparison fail even when the chosen backward convention is intentional. Passing verifies these inputs and tolerances, not every possible input.

def my_op(x):
    return (x ** 3).sum() / (1 + x.abs().sum())

x64 = torch.tensor([-1.5, -0.5, 0.7, 2.0], dtype=torch.float64, requires_grad=True)
print(torch.autograd.gradcheck(my_op, (x64,), eps=1e-6, atol=1e-8))
# True

The L-layer forward and backward passes are written out in Deep Neural Networks: Architecture and Backpropagation.

Moving to the GPU

For matrix multiplication and typical neural-network operations, keep model parameters and participating tensors on compatible devices and dtypes. The example chooses CUDA if available, then Apple MPS, then CPU. The CPU fallback also lets the example run without a supported accelerator. The simple nn.Linear module created earlier stores a trainable weight and bias; its input and output have two features per example.

device = ("cuda" if torch.cuda.is_available()
          else "mps" if torch.backends.mps.is_available()
          else "cpu")
model = model.to(device)
batch = batch.to(device)
with torch.no_grad():
    device_output = model(batch)
print("same device", device_output.device == batch.device)
print("output shape", tuple(device_output.shape))
# same device True
# output shape (3, 2)

For a module, model.to(device) moves its parameters and registered buffers in place and returns the module. Move it before constructing an optimizer. For a tensor, assign the result of batch.to(device): a conversion returns a new tensor, but if the requested dtype and device already match, it can return the original object. Device and dtype compatibility still need checking when a new batch arrives.

Reading a GPU scalar with .item(), or making a blocking copy to CPU, can force the host to wait for device work. Frequent logging can become expensive; measure before optimizing it. For interval loss reporting, accumulate loss.detach() on the device and convert when reporting. Accumulating graph-connected losses would retain their histories. Transfers and available operations also depend on the backend, so CUDA and MPS behavior should not be assumed identical.

The training loop this code slots into is built in Building a PyTorch Training Loop with nn.Module, Dataset, and DataLoader.

Exercises

1. Accumulation, seen. Create a leaf tensor and rebuild the expression p ** 2 three times, calling backward on each fresh result without clearing. Print the three gradients. Repeat with clearing before each backward. Explain why rebuilding the expression differs from calling backward repeatedly on one stored result.

You should get: an arithmetic progression in the first case and a constant in the second.

import torch
p = torch.tensor(2.0, requires_grad=True)
uncleared = []
for _ in range(3):
    (p ** 2).backward()
    uncleared.append(p.grad.item())
cleared = []
for _ in range(3):
    p.grad = None
    (p ** 2).backward()
    cleared.append(p.grad.item())
print(uncleared)
print(cleared)
# [4.0, 8.0, 12.0]
# [4.0, 4.0, 4.0]

2. What autograd saves. Use a small network to count tensor-save events during a tracked forward and during a forward inside no_grad. Explain what this tells you about backward storage, and why it is not a measurement of peak process or GPU memory.

The tracked forward saves tensors for backward; the no-grad forward does not. Counted tensor bytes can include storage already held elsewhere.

import torch
import torch.nn as nn
net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 1))
inputs = torch.ones(3, 4)
saved = []
def pack(tensor):
    saved.append((tuple(tensor.shape), tensor.numel() * tensor.element_size()))
    return tensor

def unpack(tensor):
    return tensor

with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
    output = net(inputs)
print("tracked output", output.grad_fn is not None)
print("saved for backward", len(saved) > 0)
del output
saved.clear()
with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
    with torch.no_grad():
        output = net(inputs)
print("no-grad output", output.grad_fn is None)
print("no-grad save count", len(saved))
# tracked output True
# saved for backward True
# no-grad output True
# no-grad save count 0

3. Multi-path accumulation in autograd. Define \(w = x\cdot y + x^2\) with \(x=3,y=4\), call backward(), and verify x.grad by hand. Then split \(x\) into two separate leaf tensors with the same value and confirm their gradients sum to the original.

You should get: a single gradient of 10, and two gradients of 4 and 6.

import torch
x = torch.tensor(3.0, requires_grad=True)
y = torch.tensor(4.0, requires_grad=True)
(x * y + x ** 2).backward()
print(x.grad.item(), y.grad.item())

a = torch.tensor(3.0, requires_grad=True)
# Independent leaves separate the product and square paths.
b = torch.tensor(3.0, requires_grad=True)
(a * y.detach() + b ** 2).backward()
print(a.grad.item(), b.grad.item())
# 10.0 3.0
# 4.0 6.0

Implementation details: PyTorch autograd mechanics, Tensor.detach, and Tensor.to.