Fix bug in truncated_graph_inputs

It could return duplicated truncated inputs before the changes, as well as return wrong outputs based on the nodes input order

Author: ricardoV94

pytensor/graph/basic.py

@@ -1056,47 +1056,48 @@ def truncated_graph_inputs(
                 truncated_inputs.append(node)
         # no more actions are needed
         return truncated_inputs
+
     blockers: Set[Variable] = set(ancestors_to_include)
     # enforce O(1) check for node in ancestors to include
     ancestors_to_include = blockers.copy()
 
     while candidates:
         # on any new candidate
         node = candidates.pop()
-        # check if the node is independent, never go above blockers
+
+        # There was a repeated reference to this node, we have already investigated it
+        if node in truncated_inputs:
+            continue
+
+        # check if the node is independent, never go above blockers;
         # blockers are independent nodes and ancestors to include
         if node in ancestors_to_include:
             # The case where node is in ancestors to include so we check if it depends on others
             # it should be removed from the blockers to check against the rest
-            dependent = variable_depends_on(node, blockers - {node})
+            dependent = variable_depends_on(node, ancestors_to_include - {node})
             # ancestors to include that are present in the graph (not disconnected)
             # should be added to truncated_inputs
             truncated_inputs.append(node)
             if dependent:
-                # if the ancestors to include is still dependent we need to go above,
-                # the search is not yet finished
-                # the node _has_ to have owner to be dependent
-                # so we do not check it
-                # and populate search to go above
+                # if the ancestors to include is still dependent we need to go above, the search is not yet finished
                 # owner can never be None for a dependent node
                 candidates.extend(node.owner.inputs)
         else:
             # A regular node to check
             dependent = variable_depends_on(node, blockers)
-            # all regular nodes fall to blockes
+            # all regular nodes fall to blockers
             # 1. it is dependent - further search irrelevant
             # 2. it is independent - the search node is inside the closure
             blockers.add(node)
             # if we've found an independent node and it is not in blockers so far
-            # it is a new indepenent node not present in ancestors to include
-            if not dependent:
-                # we've found an independent node
-                # do not search beyond
-                truncated_inputs.append(node)
-            else:
-                # populate search otherwise
+            # it is a new independent node not present in ancestors to include
+            if dependent:
+                # populate search if it's not an independent node
                 # owner can never be None for a dependent node
                 candidates.extend(node.owner.inputs)
+            else:
+                # otherwise, do not search beyond
+                truncated_inputs.append(node)
     return truncated_inputs
 
 

tests/graph/test_basic.py

@@ -697,55 +697,95 @@ def test_variable_depends_on():
     assert variable_depends_on(y, [y])
 
 
-def test_truncated_graph_inputs():
-    """
-    * No conditions
-        n - n - (o)
-
-    * One condition
-        n - (c) - o
-
-    * Two conditions where on depends on another, both returned
-        (c) - (c) - o
-
-    * Additional nodes are present
-           (c) - n - o
-        n - (n) -'
+class TestTruncatedGraphInputs:
+    def test_basic(self):
+        """
+        * No conditions
+            n - n - (o)
+
+        * One condition
+            n - (c) - o
+
+        * Two conditions where on depends on another, both returned
+            (c) - (c) - o
+
+        * Additional nodes are present
+               (c) - n - o
+            n - (n) -'
+
+        * Disconnected condition not returned
+            (c) - n - o
+             c
+
+        * Disconnected output is present and returned
+            (c) - (c) - o
+            (o)
+
+        * Condition on itself adds itself
+            n - (c) - (o/c)
+        """
+        x = MyVariable(1)
+        x.name = "x"
+        y = MyVariable(1)
+        y.name = "y"
+        z = MyVariable(1)
+        z.name = "z"
+        x2 = MyOp(x)
+        x2.name = "x2"
+        y2 = MyOp(y, x2)
+        y2.name = "y2"
+        o = MyOp(y2)
+        o2 = MyOp(o)
+        # No conditions
+        assert truncated_graph_inputs([o]) == [o]
+        # One condition
+        assert truncated_graph_inputs([o2], [y2]) == [y2]
+        # Condition on itself adds itself
+        assert truncated_graph_inputs([o], [y2, o]) == [o, y2]
+        # Two conditions where on depends on another, both returned
+        assert truncated_graph_inputs([o2], [y2, o]) == [o, y2]
+        # Additional nodes are present
+        assert truncated_graph_inputs([o], [y]) == [x2, y]
+        # Disconnected condition
+        assert truncated_graph_inputs([o2], [y2, z]) == [y2]
+        # Disconnected output is present
+        assert truncated_graph_inputs([o2, z], [y2]) == [z, y2]
+
+    def test_repeated_input(self):
+        """Test that truncated_graph_inputs does not return repeated inputs."""
+        x = MyVariable(1)
+        x.name = "x"
+        y = MyVariable(1)
+        y.name = "y"
+
+        trunc_inp1 = MyOp(x, y)
+        trunc_inp1.name = "trunc_inp1"
+
+        trunc_inp2 = MyOp(x, y)
+        trunc_inp2.name = "trunc_inp2"
+
+        o = MyOp(trunc_inp1, trunc_inp1, trunc_inp2, trunc_inp2)
+        o.name = "o"
+
+        assert truncated_graph_inputs([o], [trunc_inp1]) == [trunc_inp2, trunc_inp1]
+
+    def test_repeated_nested_input(self):
+        """Test that truncated_graph_inputs does not return repeated inputs."""
+        x = MyVariable(1)
+        x.name = "x"
+        y = MyVariable(1)
+        y.name = "y"
+
+        trunc_inp = MyOp(x, y)
+        trunc_inp.name = "trunc_inp"
+
+        o1 = MyOp(trunc_inp, trunc_inp, x, x)
+        o1.name = "o1"
 
-    * Disconnected condition not returned
-        (c) - n - o
-         c
+        assert truncated_graph_inputs([o1], [trunc_inp]) == [x, trunc_inp]
 
-    * Disconnected output is present and returned
-        (c) - (c) - o
-        (o)
+        # Reverse order of inputs
+        o2 = MyOp(x, x, trunc_inp, trunc_inp)
+        o2.name = "o2"
 
-    * Condition on itself adds itself
-        n - (c) - (o/c)
-    """
-    x = MyVariable(1)
-    x.name = "x"
-    y = MyVariable(1)
-    y.name = "y"
-    z = MyVariable(1)
-    z.name = "z"
-    x2 = MyOp(x)
-    x2.name = "x2"
-    y2 = MyOp(y, x2)
-    y2.name = "y2"
-    o = MyOp(y2)
-    o2 = MyOp(o)
-    # No conditions
-    assert truncated_graph_inputs([o]) == [o]
-    # One condition
-    assert truncated_graph_inputs([o2], [y2]) == [y2]
-    # Condition on itself adds itself
-    assert truncated_graph_inputs([o], [y2, o]) == [o, y2]
-    # Two conditions where on depends on another, both returned
-    assert truncated_graph_inputs([o2], [y2, o]) == [o, y2]
-    # Additional nodes are present
-    assert truncated_graph_inputs([o], [y]) == [x2, y]
-    # Disconnected condition
-    assert truncated_graph_inputs([o2], [y2, z]) == [y2]
-    # Disconnected output is present
-    assert truncated_graph_inputs([o2, z], [y2]) == [z, y2]
+        assert truncated_graph_inputs([o2], [trunc_inp]) == [trunc_inp, x]