diff --git a/CHANGELOG.md b/CHANGELOG.md index 440d066d..25aeae98 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/torchjd/autojac/_accumulation.py b/src/torchjd/autojac/_accumulation.py index 84e0702a..4cb28dc2 100644 --- a/src/torchjd/autojac/_accumulation.py +++ b/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: diff --git a/tests/unit/autojac/test_backward.py b/tests/unit/autojac/test_backward.py index 893ad1a3..178a772e 100644 --- a/tests/unit/autojac/test_backward.py +++ b/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]) diff --git a/tests/unit/autojac/test_mtl_backward.py b/tests/unit/autojac/test_mtl_backward.py index 06dca9ed..a4370178 100644 --- a/tests/unit/autojac/test_mtl_backward.py +++ b/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