Backward from a vector

y = x * 2          # a vector
y.backward()

RuntimeError: grad can be implicitly created only for scalar outputs

Why

A gradient is the derivative of one number with respect to the inputs. y is three numbers, so “the gradient of y“ is not a thing: there are three gradients, and what you get back has to be a combination of them.

backward() on a scalar works because there is only one thing it could mean, and it quietly supplies a seed of 1.0.

Saying what you meant

y.backward(torch.ones_like(y))     # the sum of the outputs
y.sum().backward()                 # identical

The argument to backward is the vector of dL/dy values that the chain rule needs. Passing ones says “the loss is the sum of these outputs”. Passing weights says “the loss is that weighted sum”.

This is why loss functions reduce. nn.MSELoss() defaults to reduction="mean" so its output is a scalar and backward() has something to start from. Passing reduction="none" gives you per-example losses and you then have to reduce them yourself before calling backward.

What is actually happening

Autograd computes a vector-Jacobian product: v @ J, where v is the seed you pass. It never builds the Jacobian, which is why it scales. Asking for the full Jacobian is a different function:

torch.autograd.functional.jacobian(f, x)

and it costs one backward pass per output element, which is why you should want the vector-Jacobian product almost always.

Your task

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

With y = x ** 2, return the gradient with respect to x of the weighted sum sum(weights * y), using a seed rather than reducing first.