fix(autojac): Materialize expanded jacs and grads

CHANGELOG.md

@@ -23,6 +23,10 @@ changelog does not include internal changes that do not affect the user.
 - Fixed `jac_to_grad` failing with an `AttributeError` (or, with `retain_jac=True`, accumulating
   a wrong gradient) when a tensor is repeated in `tensors`. Repeated tensors are now only
   considered once.
+- Fixed `backward` and `mtl_backward` storing expanded tensors in the `.jac` and `.grad` fields
+  (e.g. when a parameter is directly summed into a loss). Calling them again on the same
+  parameters, or modifying these fields in-place (e.g. with `clip_grad_norm_`), then failed with a
+  `RuntimeError`. Such tensors are now materialized before being stored.
 
 ## [0.18.1] - 2026-11-02
 

src/torchjd/autojac/_accumulation.py

@@ -1,6 +1,7 @@
 from collections.abc import Iterable
 from typing import TypeGuard
 
+import torch
 from torch import Tensor
 
 
@@ -41,7 +42,7 @@ def accumulate_jacs(params: Iterable[Tensor], jacobians: Iterable[Tensor]) -> No
             #
             # We do not detach from the computation graph because the value can have grad_fn
             # that we want to keep track of (in case it was obtained via create_graph=True).
-            param.__setattr__("jac", jac)
+            param.__setattr__("jac", _materialize_if_expanded(jac))
 
 
 def accumulate_grads(params: Iterable[Tensor], gradients: Iterable[Tensor]) -> None:
@@ -50,7 +51,15 @@ def accumulate_grads(params: Iterable[Tensor], gradients: Iterable[Tensor]) -> N
         if hasattr(param, "grad") and param.grad is not None:
             param.grad += grad
         else:
-            param.grad = grad
+            param.grad = _materialize_if_expanded(grad)
+
+
+def _materialize_if_expanded(tensor: Tensor) -> Tensor:
+    if tensor.layout == torch.strided and any(
+        stride == 0 and size > 1 for stride, size in zip(tensor.stride(), tensor.shape, strict=True)
+    ):
+        return tensor.contiguous()
+    return tensor
 
 
 def _check_expects_grad(tensor: Tensor, field_name: str) -> None:

tests/unit/autojac/test_backward.py

@@ -55,6 +55,21 @@ def test_jac_is_populated() -> None:
         assert_has_jac(a)
 
 
+def test_repeated_calls_with_expanded_jacobian() -> None:
+    """
+    Tests that backward correctly accumulates the Jacobians across several calls when the Jacobian
+    of the tensor with respect to the input is an expanded tensor (e.g. because the input is
+    directly summed).
+    """
+
+    a = tensor_([1.0, 2.0], requires_grad=True)
+
+    backward(a.sum())
+    backward(2.0 * a.sum())
+
+    assert_jac_close(a, tensor_([[3.0, 3.0]]))
+
+
 @mark.parametrize("shape", [(1, 3), (2, 3), (2, 6), (5, 8), (20, 55)])
 @mark.parametrize("manually_specify_inputs", [True, False])
 @mark.parametrize("chunk_size", [1, 2, None])

tests/unit/autojac/test_mtl_backward.py

@@ -144,6 +144,28 @@ def test_single_task() -> None:
     assert_has_grad(p1)
 
 
+def test_repeated_calls_with_expanded_task_gradient() -> None:
+    """
+    Tests that mtl_backward correctly accumulates the gradients of the task-specific parameters
+    across several calls when these gradients are expanded tensors (e.g. because the task-specific
+    parameters are directly summed).
+    """
+
+    p0 = tensor_([1.0, 2.0], requires_grad=True)
+    p1 = tensor_([3.0, 4.0], requires_grad=True)
+    p2 = tensor_([5.0, 6.0], requires_grad=True)
+
+    for _ in range(2):
+        f = tensor_([-1.0, 1.0]) @ p0
+        y1 = f + p1.sum()
+        y2 = 2.0 * f + p2.sum()
+        mtl_backward([y1, y2], features=f)
+
+    assert_jac_close(p0, tensor_([[-2.0, 2.0], [-4.0, 4.0]]))
+    assert_grad_close(p1, tensor_([2.0, 2.0]))
+    assert_grad_close(p2, tensor_([2.0, 2.0]))
+
+
 def test_incoherent_task_number_fails() -> None:
     """
     Tests that mtl_backward raises an error when called with the number of tasks losses different