Skip to content

← Autograd Mechanics step 3 of 8

Medium Primitives

Backward from a vector

Jump to the task

For an output with several elements, backward() needs an incoming gradient that says how those elements contribute to the quantity being differentiated.

import torch

x = torch.tensor([1., 2.], requires_grad=True)
y = x.square()
weights = torch.tensor([0.5, 2.])
y.backward(weights)
print(x.grad.tolist())  # [1.0, 8.0]

Here the scalar objective is 0.5 * x[0]**2 + 2 * x[1]**2. The incoming gradient with respect to y is [0.5, 2.], so the gradient with respect to x is 2 * x * weights.

Why an argument is needed

A vector-valued function has a Jacobian containing all output derivatives. Backward computes a vector-Jacobian product, combining those derivatives with the incoming gradient you supply instead of constructing the full matrix.

For a real output containing exactly one element, PyTorch can supply an incoming gradient of one automatically. This includes a tensor with shape (1,), not just a zero-dimensional scalar. An output with several elements needs an explicit gradient of matching shape, or a reduction first.

For example, y.backward(torch.ones_like(y)) and y.sum().backward() give the same input gradient when run on equivalent fresh graphs. A weighted seed likewise matches a weighted sum when the weights are fixed with respect to x.

Loss reductions

MSELoss(reduction="mean") returns a scalar. With reduction="none", it returns elementwise losses, which are not necessarily one value per example. You can reduce them or supply an appropriate incoming gradient. Computing a full Jacobian is a separate operation; its cost depends on input/output sizes and the differentiation strategy.

Your task

def weighted_grad(x: torch.Tensor, weights: torch.Tensor) -> torch.Tensor

The harness supplies a fresh floating-point leaf tensor x with gradient tracking enabled and a fixed weights tensor of the same shape. With y = x ** 2, return the gradient of sum(weights * y) with respect to x. Use weights as the incoming gradient to backward rather than first reducing the output to a scalar. Return a tensor with the same shape as x.

The starter omits that incoming gradient. It raises for the multi-element cases; a single-element case can run but does not apply arbitrary weights.