Per-Sample Gradients with vmap

Compute one gradient vector per training example, using the same weights for every example. Keep those gradients separate instead of averaging them.

Your task

Implement per_sample_grads(w, X, y) by composing automatic differentiation with vmap. The floating-point inputs share a dtype and device:

• w has shape (d,) and contains the shared weights.
• X has shape (batch, d), with one example per row.
• y has shape (batch,), with one scalar target per example.

For example i, the loss is (dot(w, X[i]) - y[i])**2. Return its gradient with respect to w in row i of a (batch, d) tensor. Preserve dtype and device and leave all inputs unchanged. Do not update the weights.

There is no factor of 1/2 and no division by batch size in this loss:

residual[i] = dot(w, X[i]) - y[i]
gradient[i, j] = 2 * residual[i] * X[i, j]

For w=[1, 0], X=[[2, 3], [4, 5]] and y=[0, 0], residuals are [2, 4] and the result is [[8, 12], [32, 40]]. Returning their sum or mean would lose the separate gradient for each example.

Share the weights, map the examples

Define a scalar loss_fn(w, xi, yi) for one example. Differentiate its first argument, w, then map that gradient function across X and y:

• PyTorch: use torch.func.grad(loss_fn) and torch.func.vmap(grad_fn, in_dims=(None, 0, 0))(w, X, y).
• JAX: use jax.grad(loss_fn, argnums=0) and jax.vmap(grad_fn, in_axes=(None, 0, 0))(w, X, y).

None means the whole weight vector is shared; 0 selects one example and target along their leading axes. You do not need to expand w into repeated rows. This is vectorized computation, not a request to distribute work across devices.

In the diagram, each inner call receives a (3,) weight vector, a (3,) example and a scalar target. Its (3,) gradient becomes one row of (8, 3). [] denotes a scalar shape. The controls illustrate axis placement rather than checking whether arbitrary choices make a valid loss function.