Add dprint shortcut to FunctionGraph and Function

Author: ricardoV94

pytensor/compile/function/types.py

@@ -1097,6 +1097,18 @@ def sync_shared(self):
         # NOTE: sync was needed on old gpu backend
         pass
 
+    def dprint(self, **kwargs):
+        """Debug print itself
+
+        Parameters
+        ----------
+        kwargs:
+            Optional keyword arguments to pass to debugprint function.
+        """
+        from pytensor.printing import debugprint
+
+        return debugprint(self, **kwargs)
+
 
 # pickling/deepcopy support for Function
 def _pickle_Function(f):

pytensor/graph/fg.py

@@ -927,3 +927,15 @@ def __contains__(self, item: Variable | Apply) -> bool:
             return item in self.apply_nodes
         else:
             raise TypeError()
+
+    def dprint(self, **kwargs):
+        """Debug print itself
+
+        Parameters
+        ----------
+        kwargs:
+            Optional keyword arguments to pass to debugprint function.
+        """
+        from pytensor.printing import debugprint
+
+        return debugprint(self, **kwargs)

tests/compile/function/test_types.py

@@ -16,6 +16,7 @@
 from pytensor.graph.rewriting.basic import OpKeyGraphRewriter, PatternNodeRewriter
 from pytensor.graph.utils import MissingInputError
 from pytensor.link.vm import VMLinker
+from pytensor.printing import debugprint
 from pytensor.tensor.math import dot, tanh
 from pytensor.tensor.math import sum as pt_sum
 from pytensor.tensor.type import (
@@ -862,6 +863,12 @@ def test_key_string_requirement(self):
         with pytest.raises(AssertionError):
             function([x], outputs={(1, "b"): x, 1.0: x**2})
 
+    def test_dprint(self):
+        x = pt.scalar("x")
+        out = x + 1
+        f = function([x], out)
+        assert f.dprint(file="str") == debugprint(f, file="str")
+
 
 class TestPicklefunction:
     def test_deepcopy(self):

tests/graph/test_fg.py

@@ -8,6 +8,7 @@
 from pytensor.graph.basic import NominalVariable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.utils import MissingInputError
+from pytensor.printing import debugprint
 from tests.graph.utils import (
     MyConstant,
     MyOp,
@@ -706,3 +707,9 @@ def test_nominals(self):
         assert nm2 not in fg.inputs
         assert nm in fg.variables
         assert nm2 in fg.variables
+
+    def test_dprint(self):
+        r1, r2 = MyVariable("x"), MyVariable("y")
+        o1 = op1(r1, r2)
+        fg = FunctionGraph([r1, r2], [o1], clone=False)
+        assert fg.dprint(file="str") == debugprint(fg, file="str")