fix(autojac): Handle repeated tensors in jac_to_grad

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
 

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 "

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."""