Add helper to extract variable graph_input

Author: Dhruvanshu-Joshi

pytensor/graph/basic.py

@@ -907,6 +907,35 @@ def graph_inputs(
     yield from (r for r in ancestors(graphs, blockers) if r.owner is None)
 
 
+def required_graph_inputs(graph):
+    """
+    Get the inputs into PyTensor variables
+
+    Parameters
+    ----------
+        graph: PyTensor `Variable` instances
+        Output `Variable` instances from which to search backward through
+        owners.
+
+    Returns
+    -------
+        List of tensor variables that are input nodes with no owner, in the order
+        found by a left-recursive depth-first search started at the nodes in `graphs`.
+    """
+    from pytensor.compile.sharedvalue import SharedVariable
+
+    if isinstance(graph, tuple | list):
+        graph = graph
+    else:
+        graph = [graph]
+
+    return [
+        v
+        for v in graph_inputs(graph)
+        if isinstance(v, Variable) and not isinstance(v, Constant | SharedVariable)
+    ]
+
+
 def vars_between(
     ins: Collection[Variable], outs: Iterable[Variable]
 ) -> Generator[Variable, None, None]:

tests/graph/test_basic.py

@@ -8,6 +8,7 @@
 from pytensor import tensor as pt
 from pytensor.graph.basic import (
     Apply,
+    Constant,
     NominalVariable,
     Variable,
     ancestors,
@@ -23,6 +24,7 @@
     io_toposort,
     list_of_nodes,
     orphans_between,
+    required_graph_inputs,
     truncated_graph_inputs,
     variable_depends_on,
     vars_between,
@@ -61,6 +63,16 @@ def MyVariable(thingy):
     return Variable(MyType(thingy), None, None)
 
 
+def MyConstant(name, thingy, data=None):
+    return Constant(MyType(thingy), data, name=name)
+
+
+# def MyConstant(thingy):
+#     return Constant(MyType(thingy), None, None)
+# def MyVariable(thingy):
+#     return Variable(MyType(thingy), None, None)
+
+
 class MyOp(Op):
     __props__ = ()
 
@@ -75,6 +87,20 @@ def perform(self, *args, **kwargs):
         raise NotImplementedError("No Python implementation available.")
 
 
+class MyCustomOp(Op):
+    __props__ = ()
+
+    def make_node(self, *inputs):
+        for input in inputs:
+            assert isinstance(input, Variable)
+            assert isinstance(input.type, MyType)
+        outputs = [MyVariable(sum(input.type.thingy for input in inputs))]
+        return Apply(self, list(inputs), outputs)
+
+    def perform(self, *args, **kwargs):
+        raise NotImplementedError("No Python implementation available.")
+
+
 MyOp = MyOp()
 
 
@@ -501,6 +527,16 @@ def test_graph_inputs():
     assert res_list == [r3, r1, r2]
 
 
+def test_required_graph_inputs():
+    x = pt.fscalar()
+    y = pt.constant(2)
+    z = shared(1)
+    a = pt.sum(x + y + z)
+
+    res = list(required_graph_inputs([a]))
+    assert res == [x]
+
+
 def test_variables_and_orphans():
     r1, r2, r3 = MyVariable(1), MyVariable(2), MyVariable(3)
     o1 = MyOp(r1, r2)