Simplify _get_leaf_tensors (#209)

* Change it to only work with tensors and excluded that have a grad_fn
* Stop making it exclude the excluded tensors themselves (only exclude their grad_fn from the graph traversal)
* Remove tests that expected the leaves to be excluded
* Add tests to verify that errors are raised when necessary

Author: ValerianRey

src/torchjd/autojac/_utils.py

@@ -37,18 +37,26 @@ def _get_leaf_tensors(tensors: Iterable[Tensor], excluded: Iterable[Tensor]) ->
     """
     Gets the leaves of the autograd graph of all specified ``tensors``.
 
-    :param tensors: Tensors from which the graph traversal should start.
-    :param excluded: Tensors that should be excluded from the results and whose grad_fn should be
-        excluded from the graph traversal.
+    :param tensors: Tensors from which the graph traversal should start. They should all require
+        grad and not be leaves.
+    :param excluded: Tensors whose grad_fn should be excluded from the graph traversal. They should
+        all require grad and not be leaves.
+
     """
 
+    if any([tensor.grad_fn is None for tensor in tensors]):
+        raise ValueError("All `tensors` should have a `grad_fn`.")
+
+    if any([tensor.grad_fn is None for tensor in excluded]):
+        raise ValueError("All `excluded` tensors should have a `grad_fn`.")
+
     accumulate_grads = _get_descendant_accumulate_grads(
-        roots={tensor.grad_fn for tensor in tensors if tensor.grad_fn is not None},
+        roots={tensor.grad_fn for tensor in tensors},
         excluded_nodes={tensor.grad_fn for tensor in excluded},
     )
     leaves = {g.variable for g in accumulate_grads}
 
-    return leaves - set(excluded)
+    return leaves
 
 
 def _get_descendant_accumulate_grads(roots: set[Node], excluded_nodes: set[Node]) -> set[Node]:

tests/unit/autojac/test_utils.py

@@ -1,5 +1,5 @@
 import torch
-from pytest import mark
+from pytest import mark, raises
 from torch.nn import Linear, MSELoss, ReLU, Sequential
 from unit.conftest import DEVICE
 
@@ -61,24 +61,6 @@ def test_get_leaf_tensors_excluded_2():
     assert leaves == {p1, p2}
 
 
-def test_get_leaf_tensors_excluded_3():
-    """
-    Tests that _get_leaf_tensors works correctly when some tensors are excluded from the search.
-
-    In this case, one of the leaves itself is excluded.
-    """
-
-    p1 = torch.tensor([1.0, 2.0], requires_grad=True, device=DEVICE)
-    p2 = torch.tensor([3.0, 4.0], requires_grad=True, device=DEVICE)
-    p3 = torch.tensor([5.0, 6.0], requires_grad=True, device=DEVICE)
-
-    y1 = torch.tensor([-1.0, 1.0], device=DEVICE) @ p1 + p2.sum()
-    y2 = (p1**2).sum() + p2.norm() + p3.sum()
-
-    leaves = _get_leaf_tensors(tensors=[y1, y2], excluded={p3})
-    assert leaves == {p1, p2}
-
-
 def test_get_leaf_tensors_leaf_not_requiring_grad():
     """
     Tests that _get_leaf_tensors does not include tensors that do not require grad in its results.
@@ -113,25 +95,6 @@ def test_get_leaf_tensors_model():
     assert leaves == set(model.parameters())
 
 
-def test_get_leaf_tensors_model_excluded_1():
-    """
-    Tests that _get_leaf_tensors works correctly when the autograd graph is generated by a simple
-    sequential model, and some of the model's parameters are excluded.
-    """
-
-    x = torch.randn(16, 10)
-    y = torch.randn(16, 1)
-
-    model = Sequential(Linear(10, 5), ReLU(), Linear(5, 1))
-    loss_fn = MSELoss(reduction="none")
-
-    y_hat = model(x)
-    losses = loss_fn(y_hat, y)
-
-    leaves = _get_leaf_tensors(tensors=[losses], excluded=set(model[0].parameters()))
-    assert leaves == set(model[2].parameters())
-
-
 def test_get_leaf_tensors_model_excluded_2():
     """
     Tests that _get_leaf_tensors works correctly when the autograd graph is generated by a simple
@@ -197,11 +160,38 @@ def test_get_leaf_tensors_deep(depth: int):
 
 
 def test_get_leaf_tensors_leaf():
+    """Tests that _get_leaf_tensors raises an error some of the provided tensors are leaves."""
+
+    a = torch.tensor(1.0, requires_grad=True, device=DEVICE)
+    with raises(ValueError):
+        _ = _get_leaf_tensors(tensors=[a], excluded=set())
+
+
+def test_get_leaf_tensors_tensor_not_requiring_grad():
     """
-    Tests that _get_leaf_tensors correctly returns an empty set when the provided tensors are
-    leaves.
+    Tests that _get_leaf_tensors raises an error some of the provided tensors do not require grad.
     """
 
-    a = torch.tensor(1.0, requires_grad=True, device=DEVICE)
-    leaves = _get_leaf_tensors(tensors=[a], excluded=set())
-    assert leaves == set()
+    a = torch.tensor(1.0, requires_grad=False, device=DEVICE) * 2
+    with raises(ValueError):
+        _ = _get_leaf_tensors(tensors=[a], excluded=set())
+
+
+def test_get_leaf_tensors_excluded_leaf():
+    """Tests that _get_leaf_tensors raises an error some of the excluded tensors are leaves."""
+
+    a = torch.tensor(1.0, requires_grad=True, device=DEVICE) * 2
+    b = torch.tensor(2.0, requires_grad=True, device=DEVICE)
+    with raises(ValueError):
+        _ = _get_leaf_tensors(tensors=[a], excluded={b})
+
+
+def test_get_leaf_tensors_excluded_not_requiring_grad():
+    """
+    Tests that _get_leaf_tensors raises an error some of the excluded tensors do not require grad.
+    """
+
+    a = torch.tensor(1.0, requires_grad=True, device=DEVICE) * 2
+    b = torch.tensor(2.0, requires_grad=False, device=DEVICE) * 2
+    with raises(ValueError):
+        _ = _get_leaf_tensors(tensors=[a], excluded={b})