PyTorch Gradients of Vectors
A vector is a row of values in one tensor. backward() works on a scalar, so a vector answer is summed into a scalar first.
A vector answer fails
This call fails on purpose. It raises a RuntimeError, and the stable last line says grad can be implicitly created only for scalar outputs.
import torch
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
y.backward()Sum first, then backward()
The sum turns the row of answers into one scalar. backward() then fills one gradient per input value.
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
total = y.sum()
total.backward()
print(x.grad)Each value gets its own gradient
The gradient of v ** 3 is 3 * v ** 2 for each value. The two inputs get 48 and 75.
v = torch.tensor([4.0, 5.0], requires_grad=True)
(v ** 3).sum().backward()
print(v.grad)Example
import torch
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(5.0, requires_grad=True)
c = torch.tensor(2.0, requires_grad=True)
total = a ** 2 + 3 * b + c ** 3
total.backward()
print("Total:", total.item())
print("Gradient for a:", a.grad.item())
print("Gradient for b:", b.grad.item())
print("Gradient for c:", c.grad.item())Key points
- Scalar: one single number
- backward() on a vector answer raises a RuntimeError
- Use .sum() to turn a vector answer into a scalar
- Each input value gets its own gradient in .grad
- The gradient row lines up with the input row, value by value
- In short, sum first, then backward() fills one gradient per input
In this example, a, b, and c are three separate tracked inputs. total adds a ** 2, 3 * b, and c ** 3 into one scalar. backward() then fills one gradient per input. The prints read the values out with .item(). total is 27.0. a gets 4.0, from 2 * a at a = 2. b gets 3.0, and c gets 12.0, from 3 * c ** 2 at c = 2.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
Why is .sum() called before backward() on a vector answer?