[SajalDevX]

When a parameter goes straight into a sum (e.g. p.sum()), autograd returns its gradient as an expanded tensor (stride 0). torch.autograd.backward copies such gradients before storing them in .grad, but accumulate_jacs and accumulate_grads store them as is (since #511 removed the clone). The stored .jac / .grad then can't be modified in-place, so a second backward / mtl_backward call on the same parameters (gradient accumulation) or a clip_grad_norm_ fails:

import torch
from torchjd.autojac import mtl_backward

shared = torch.nn.Linear(3, 4)
b = torch.nn.Parameter(torch.zeros(4))
for _ in range(2):
    f = shared(torch.randn(5, 3))
    mtl_backward([f.pow(2).mean(), (f.sum(0) + b).sum()], features=f)
# RuntimeError: unsupported operation: more than one element of the written-to tensor refers to a single memory location.

Same thing with backward(a.sum()); backward(2.0 * a.sum()).

This PR makes accumulate_jacs and accumulate_grads call .contiguous() on the value before storing it, only when it has a stride-0 dimension of size > 1. Other tensors (including sparse ones) are still stored without a copy.

Tested:

• Added test_repeated_calls_with_expanded_jacobian (test_backward.py) and test_repeated_calls_with_expanded_task_gradient (test_mtl_backward.py). Both fail on main with the RuntimeError above and pass with the fix.
• uv run pytest tests/unit -W error (CPU): 3714 passed, 66 skipped, 33 xfailed.
• uv run ruff check, uv run ruff format, uv run ty check (only the existing unresolved-import errors for the optional plotting deps).

[ValerianRey]

Interesting. IIUC this fixes a quite niche bug (multiple accumulations + expanded grad / jac) at the expense of extra memory in some cases.

Could you move the fix to the case where we do param.jac += jac? What I mean is, wait for a second accumulation to be done, and when it happens (i.e. a jac field is already there), materialize the existing grad if its has some zero stride, and then only do the accumulation.

That would be more memory efficient in most cases. But maybe your current fix also fixes other issues that can trigger even with a single accumulation. So idk what's best.

[SajalDevX]

Good point on memory. It can also break after a single call though: one mtl_backward followed by clip_grad_norm_ crashes, since clipping scales .grad in place. Would it work to keep the copy for .grad, which matches what torch's own backward stores, and make .jac lazy as you suggest, only materializing when a second jac is accumulated?