Remove parameter excluded from get_leaf_edges

* It was never used
* Since roots was actually not modified in-place, this also does excluded = roots, which makes roots actually be modified in place (like the documentation says). So it behave differently than before but in an expected way.

Author: ValerianRey

src/torchjd/autogram/_edge_registry.py

@@ -23,25 +23,21 @@ def register(self, edge: GradientEdge) -> None:
         """
         self._edges.add(edge)
 
-    def get_leaf_edges(
-        self, roots: set[GradientEdge], excluded: set[GradientEdge]
-    ) -> set[GradientEdge]:
+    def get_leaf_edges(self, roots: set[GradientEdge]) -> set[GradientEdge]:
         """
         Compute a minimal subset of edges that yields the same differentiation graph traversal: the
         leaf edges. Specifically, this removes edges that are reachable from other edges in the
         differentiation graph, avoiding the need to keep gradients in memory for all edges
         simultaneously.
 
         :param roots: Roots of the graph traversal. Modified in-place.
-        :param excluded: GradientEdges that stop graph traversal. Modified in-place.
         :returns: Minimal subset of leaf edges.
         """
 
-        roots.difference_update(excluded)
         nodes_to_traverse = deque((child, root) for root in roots for child in _next_edges(root))
         result = {root for root in roots if root in self._edges}
 
-        excluded.update(roots)
+        excluded = roots
         while nodes_to_traverse:
             node, origin = nodes_to_traverse.popleft()
             if node in self._edges:

src/torchjd/autogram/_engine.py

@@ -175,7 +175,7 @@ def compute_gramian(self, output: Tensor) -> Tensor:
     def _compute_square_gramian(self, output: Tensor) -> Tensor:
         self._module_hook_manager.gramian_accumulation_phase = True
 
-        leaf_targets = list(self._target_edges.get_leaf_edges({get_gradient_edge(output)}, set()))
+        leaf_targets = list(self._target_edges.get_leaf_edges({get_gradient_edge(output)}))
 
         def differentiation(_grad_output: Tensor) -> tuple[Tensor, ...]:
             return torch.autograd.grad(

tests/unit/autogram/test_edge_registry.py

@@ -18,7 +18,7 @@ def test_all_edges_are_leaves1():
         edge_registry.register(get_gradient_edge(tensor))
 
     expected_leaves = {get_gradient_edge(tensor) for tensor in [a, b, c]}
-    leaves = edge_registry.get_leaf_edges({get_gradient_edge(d)}, set())
+    leaves = edge_registry.get_leaf_edges({get_gradient_edge(d)})
     assert leaves == expected_leaves
 
 
@@ -42,7 +42,7 @@ def test_all_edges_are_leaves2():
         edge_registry.register(get_gradient_edge(tensor))
 
     expected_leaves = {get_gradient_edge(tensor) for tensor in [e, g]}
-    leaves = edge_registry.get_leaf_edges({get_gradient_edge(e), get_gradient_edge(g)}, set())
+    leaves = edge_registry.get_leaf_edges({get_gradient_edge(e), get_gradient_edge(g)})
     assert leaves == expected_leaves
 
 
@@ -63,7 +63,7 @@ def test_some_edges_are_not_leaves1():
         edge_registry.register(get_gradient_edge(tensor))
 
     expected_leaves = {get_gradient_edge(tensor) for tensor in [a, b, c, d]}
-    leaves = edge_registry.get_leaf_edges({get_gradient_edge(g)}, set())
+    leaves = edge_registry.get_leaf_edges({get_gradient_edge(g)})
     assert leaves == expected_leaves
 
 
@@ -88,5 +88,5 @@ def test_some_edges_are_not_leaves2():
         edge_registry.register(get_gradient_edge(tensor))
 
     expected_leaves = {get_gradient_edge(tensor) for tensor in [a, c, d]}
-    leaves = edge_registry.get_leaf_edges({get_gradient_edge(g)}, set())
+    leaves = edge_registry.get_leaf_edges({get_gradient_edge(g)})
     assert leaves == expected_leaves