diff --git a/CHANGELOG.md b/CHANGELOG.md index 14286b4f..440d066d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,9 @@ changelog does not include internal changes that do not affect the user. `initial_weights` and `leak`). These tensors are now registered as non-persistent buffers, so they follow the dtype and device of the aggregator or weighting without being added to its `state_dict`. +- 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. ## [0.18.1] - 2026-11-02 diff --git a/src/torchjd/autojac/_jac_to_grad.py b/src/torchjd/autojac/_jac_to_grad.py index 8ceb5641..b5a10b13 100644 --- a/src/torchjd/autojac/_jac_to_grad.py +++ b/src/torchjd/autojac/_jac_to_grad.py @@ -15,6 +15,7 @@ from torchjd.linalg import Matrix, PSDMatrix from ._accumulation import TensorWithJac, accumulate_grads, is_tensor_with_jac +from ._transform import OrderedSet from ._utils import check_consistent_first_dimension @@ -62,7 +63,8 @@ def jac_to_grad( into their ``.grad`` fields. :param tensors: The tensors whose ``.jac`` fields should be aggregated. All Jacobians must - have the same first dimension (e.g. number of losses). + have the same first dimension (e.g. number of losses). Repeated tensors are only considered + once. :param aggregator: The aggregator used to reduce the Jacobians into gradients. If it uses a :class:`~torchjd.aggregation.Weighting` to combine the rows of the Jacobians, ``jac_to_grad`` will also return the computed weights. @@ -118,7 +120,7 @@ def jac_to_grad( """ tensors_ = list[TensorWithJac]() - for t in tensors: + for t in OrderedSet(tensors): if not is_tensor_with_jac(t): raise ValueError( "Some `jac` fields were not populated. Did you use `autojac.backward` or " diff --git a/tests/unit/autojac/test_jac_to_grad.py b/tests/unit/autojac/test_jac_to_grad.py index 89bde206..4a1b376f 100644 --- a/tests/unit/autojac/test_jac_to_grad.py +++ b/tests/unit/autojac/test_jac_to_grad.py @@ -91,6 +91,30 @@ def test_single_tensor() -> None: assert_grad_close(t, g) +@mark.parametrize("retain_jac", [True, False]) +@mark.parametrize("optimize", [True, False]) +def test_repeated_tensors(retain_jac: bool, optimize: bool) -> None: + """Tests that jac_to_grad considers repeated tensors only once.""" + + aggregator = base_agg() + t1 = tensor_(1.0, requires_grad=True) + t2 = tensor_([2.0, 3.0], requires_grad=True) + jac = tensor_([[-4.0, 1.0, 1.0], [6.0, 1.0, 1.0]]) + t1.__setattr__("jac", jac[:, 0]) + t2.__setattr__("jac", jac[:, 1:]) + g = aggregator(jac) + + jac_to_grad( + [t1, t2, t1], + aggregator, + retain_jac=retain_jac, + optimize_gramian_computation=optimize, + ) + + assert_grad_close(t1, g[0]) + assert_grad_close(t2, g[1:]) + + def test_no_jac_field() -> None: """Tests that jac_to_grad fails when a tensor does not have a jac field."""