Repository: TorchJD/torchjd

Issue:
`mtl_backward` should raise an error when shared and task params overlap
To solve this, we need to:
- Add a check in `mtl_backward`.
- Add a test to verify that such a call raises an error.
- Changelog entry about potentially breaking bug fix 

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/TorchJD/torchjd/pull/202?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TorchJD) Report
All modified and coverable lines are covered by tests :white_check_mark:

| [Files with missing lines](https://app.codecov.io/gh/TorchJD/torchjd/pull/202?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TorchJD) | Coverage Δ | |
|---|---|---|
| [src/torchjd/autojac/mtl\_backward.py](https://app.codecov.io/gh/TorchJD/torchjd/pull/202?src=pr&el=tree&filepath=src%2Ftorchjd%2Fautojac%2Fmtl_backward.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TorchJD#diff-c3JjL3RvcmNoamQvYXV0b2phYy9tdGxfYmFja3dhcmQucHk=) | `100.00% <100.00%> (ø)` | |


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3fb5add..07e4713 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -31,6 +31,11 @@ changes that do not affect the user.
   `pyproject.toml` (from a `[tool.pdm.dev-dependencies]` to a `[dependency-groups]` section). This
   should only affect development dependencies.
 
+### Fixed
+
+- **BREAKING**: Added a check in `mtl_backward` to ensure that `tasks_params` and `shared_params`
+  have no overlap. Previously, the behavior in this scenario was quite arbitrary.
+
 ## [0.2.2] - 2024-11-11
 
 ### Added
diff --git a/src/torchjd/autojac/mtl_backward.py b/src/torchjd/autojac/mtl_backward.py
index a900dde..0c81342 100644
--- a/src/torchjd/autojac/mtl_backward.py
+++ b/src/torchjd/autojac/mtl_backward.py
@@ -92,7 +92,7 @@ def mtl_backward(
         raise ValueError("`features` cannot be empty.")
 
     _check_retain_graph_compatible_with_chunk_size(features, retain_graph, parallel_chunk_size)
-
+    _check_no_overlap(shared_params, tasks_params)
     _check_losses_are_scalar(losses)
 
     if len(losses) == 0:
@@ -167,3 +167,12 @@ def _check_losses_are_scalar(losses: Sequence[Tensor]) -> None:
     for loss in losses:
         if loss.ndim > 0:
             raise ValueError("`losses` should contain only scalars.")
+
+
+def _check_no_overlap(shared_params: Iterable[Tensor], tasks_params: Sequence[Iterable[Tensor]]):
+    task_param_set = {param for task_params in tasks_params for param in task_params}
+    shared_param_set = set(shared_params)
+    intersection = task_param_set.intersection(shared_param_set)
+
+    if len(intersection) != 0:
+        raise ValueError("`tasks_params` should contain no tensor in common with `shared_params`.")


Test patch:
diff --git a/tests/unit/autojac/test_mtl_backward.py b/tests/unit/autojac/test_mtl_backward.py
index df76e30..af5234e 100644
--- a/tests/unit/autojac/test_mtl_backward.py
+++ b/tests/unit/autojac/test_mtl_backward.py
@@ -541,3 +541,28 @@ def test_mtl_backward_task_params_are_subset_of_other_task_params():
 
     J = torch.tensor([[-p1, p1], [-p1 * p2, p1 * p2]], device=DEVICE)
     assert_close(p0.grad, A(J))
+
+
+def test_mtl_backward_shared_params_overlap_with_tasks_params():
+    """
+    Tests that mtl_backward raises an error when the set of shared params overlaps with the set of
+    task-specific params.
+    """
+
+    p0 = torch.tensor([1.0, 2.0], requires_grad=True, device=DEVICE)
+    p1 = torch.tensor(2.0, requires_grad=True, device=DEVICE)
+    p2 = torch.tensor(3.0, requires_grad=True, device=DEVICE)
+
+    r = torch.tensor([-1.0, 1.0], device=DEVICE) @ p0
+    y1 = r * p1
+    y2 = p0.sum() * r * p2
+
+    with raises(ValueError):
+        mtl_backward(
+            losses=[y1, y2],
+            features=[r],
+            A=UPGrad(),
+            tasks_params=[[p1], [p0, p2]],  # Problem: p0 is also shared
+            shared_params=[p0],
+            retain_graph=True,
+        )