Improve test_edge_registry

- Add test_all_edges_are_leaves2
- Simplify tests

Author: ValerianRey

tests/unit/autogram/test_edge_registry.py

@@ -4,32 +4,51 @@
 from torchjd.autogram._edge_registry import EdgeRegistry
 
 
-def test_all_edges_are_leaves():
+def test_all_edges_are_leaves1():
+    """Tests that get_leaf_edges works correctly when all edges are already leaves."""
+
     a = randn_([3, 4], requires_grad=True)
     b = randn_([4], requires_grad=True)
     c = randn_([3], requires_grad=True)
 
     d = (a @ b) + c
 
     edge_registry = EdgeRegistry()
+    for tensor in [a, b, c]:
+        edge_registry.register(get_gradient_edge(tensor))
 
-    registered_edges = {
-        a: get_gradient_edge(a),
-        b: get_gradient_edge(b),
-        c: get_gradient_edge(c),
-    }
+    expected_leaves = {get_gradient_edge(tensor) for tensor in [a, b, c]}
+    leaves = edge_registry.get_leaf_edges({get_gradient_edge(d)}, set())
+    assert leaves == expected_leaves
 
-    expected_leaves = {registered_edges[a], registered_edges[b], registered_edges[c]}
 
-    for edge in registered_edges.values():
-        edge_registry.register(edge)
+def test_all_edges_are_leaves2():
+    """
+    Tests that get_leaf_edges works correctly when all edges are already leaves of the graph of
+    edges leading to them, but are not leaves of the autograd graph.
+    """
 
-    leaves = edge_registry.get_leaf_edges({get_gradient_edge(d)}, set())
+    a = randn_([3, 4], requires_grad=True)
+    b = randn_([4], requires_grad=True)
+    c = randn_([4], requires_grad=True)
+    d = randn_([4], requires_grad=True)
+
+    e = a * b
+    f = e + c
+    g = f + d
+
+    edge_registry = EdgeRegistry()
+    for tensor in [e, g]:
+        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())
     assert leaves == expected_leaves
 
 
 def test_some_edges_are_not_leaves1():
+    """Tests that get_leaf_edges works correctly when some edges are leaves and some are not."""
+
     a = randn_([3, 4], requires_grad=True)
     b = randn_([4], requires_grad=True)
     c = randn_([4], requires_grad=True)
@@ -40,33 +59,21 @@ def test_some_edges_are_not_leaves1():
     g = f + d
 
     edge_registry = EdgeRegistry()
+    for tensor in [a, b, c, d, e, f, g]:
+        edge_registry.register(get_gradient_edge(tensor))
 
-    registered_edges = {
-        a: get_gradient_edge(a),
-        b: get_gradient_edge(b),
-        c: get_gradient_edge(c),
-        d: get_gradient_edge(d),
-        e: get_gradient_edge(e),
-        f: get_gradient_edge(f),
-        g: get_gradient_edge(g),
-    }
-
-    expected_leaves = {
-        registered_edges[a],
-        registered_edges[b],
-        registered_edges[c],
-        registered_edges[d],
-    }
-
-    for edge in registered_edges.values():
-        edge_registry.register(edge)
-
+    expected_leaves = {get_gradient_edge(tensor) for tensor in [a, b, c, d]}
     leaves = edge_registry.get_leaf_edges({get_gradient_edge(g)}, set())
-
     assert leaves == expected_leaves
 
 
 def test_some_edges_are_not_leaves2():
+    """
+    Tests that get_leaf_edges works correctly when some edges are leaves and some are not. This
+    time, not all tensors in the graph are registered so not all leavese in the graph have to be
+    returned.
+    """
+
     a = randn_([3, 4], requires_grad=True)
     b = randn_([4], requires_grad=True)
     c = randn_([4], requires_grad=True)
@@ -77,24 +84,9 @@ def test_some_edges_are_not_leaves2():
     g = f + d
 
     edge_registry = EdgeRegistry()
+    for tensor in [a, c, d, e, g]:
+        edge_registry.register(get_gradient_edge(tensor))
 
-    registered_edges = {
-        a: get_gradient_edge(a),
-        c: get_gradient_edge(c),
-        d: get_gradient_edge(d),
-        e: get_gradient_edge(e),
-        g: get_gradient_edge(g),
-    }
-
-    expected_leaves = {
-        registered_edges[a],
-        registered_edges[c],
-        registered_edges[d],
-    }
-
-    for edge in registered_edges.values():
-        edge_registry.register(edge)
-
+    expected_leaves = {get_gradient_edge(tensor) for tensor in [a, c, d]}
     leaves = edge_registry.get_leaf_edges({get_gradient_edge(g)}, set())
-
     assert leaves == expected_leaves