Add option to print inner graphs in debugprint function

[Aarsh-Wankar]

I added an option print_inner_graphs to print inner graphs in the debugprint function.

Description

A boolean argument print_inner_graphs is provided to the debugprint function, which defaults to True. In case we don't want to print the inner graphs, we can set it to False.
For example:

import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import get_default_mode

n = pt.iscalar("n")
x0 = pt.vector("x0")
xs, _ = pytensor.scan(lambda xtm1: xtm1 + 1, outputs_info=[x0], n_steps=n)

mode = get_default_mode().including("scan_save_mem")
fn = pytensor.function([n, x0], xs, mode=mode, on_unused_input="ignore")
fn.dprint()

# Output:
  Subtensor{start:stop} [id A] 9
   ├─ Scan{scan_fn, while_loop=False, inplace=all} [id B] 8
   │  ├─ n [id C]
   │  └─ SetSubtensor{:stop} [id D] 7
   │     ├─ AllocEmpty{dtype='float64'} [id E] 6
   │     │  ├─ Composite{...}.2 [id F] 0
   │     │  │  └─ n [id C]
   │     │  └─ Shape_i{0} [id G] 5
   │     │     └─ x0 [id H]
   │     ├─ Unbroadcast{0} [id I] 4
   │     │  └─ ExpandDims{axis=0} [id J] 3
   │     │     └─ x0 [id H]
   │     └─ 1 [id K]
   ├─ ScalarFromTensor [id L] 2
   │  └─ Composite{...}.1 [id F] 0
   │     └─ ···
   └─ ScalarFromTensor [id M] 1
      └─ Composite{...}.0 [id F] 0
         └─ ···
  
  Inner graphs:
  
  Scan{scan_fn, while_loop=False, inplace=all} [id B]
   ← Add [id N]
      ├─ [1.] [id O]
  ...
      │     ├─ maximum [id X] 't8'
      │     │  └─ ···
      │     └─ 1 [id W]
      └─ 1 [id U]

fn.dprint(print_inner_graphs = False)

# Output:
Subtensor{start:stop} [id A] 9
 ├─ Scan{scan_fn, while_loop=False, inplace=all} [id B] 8
 │  ├─ n [id C]
 │  └─ SetSubtensor{:stop} [id D] 7
 │     ├─ AllocEmpty{dtype='float64'} [id E] 6
 │     │  ├─ Composite{...}.2 [id F] 0
 │     │  │  └─ n [id C]
 │     │  └─ Shape_i{0} [id G] 5
 │     │     └─ x0 [id H]
 │     ├─ Unbroadcast{0} [id I] 4
 │     │  └─ ExpandDims{axis=0} [id J] 3
 │     │     └─ x0 [id H]
 │     └─ 1 [id K]
 ├─ ScalarFromTensor [id L] 2
 │  └─ Composite{...}.1 [id F] 0
 │     └─ ···
 └─ ScalarFromTensor [id M] 1
    └─ Composite{...}.0 [id F] 0
       └─ ···

Related Issue

• Closes Add option not to print inner graphs in debug_print #1285
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pytensor--1293.org.readthedocs.build/en/1293/

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.03%. Comparing base (c822a8e) to head (04ebe0a).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1293   +/-   ##
=======================================
  Coverage   82.03%   82.03%           
=======================================
  Files         188      188           
  Lines       48567    48567           
  Branches     8675     8675           
=======================================
  Hits        39841    39841           
  Misses       6574     6574           
  Partials     2152     2152           

Files with missing lines	Coverage Δ
pytensor/printing.py	51.43% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ricardoV94]

There's some logic above about collecting this inner_graph_vars, can we avoid doing that work when print_inner_graphs=False?

Also wonder if we could instead specify the "depth" of the inner graphs we are interested in, like we have the depth argument now for the depth of the graph. So if inner_graphs_depth=-1, we have the default behavior (show all inner_graphs), but if 0, we don't show any, and if 1 we show one level of inner graphs, but wouldn't show inner graphs inside other inner graphs, (and so on)?.

Can you spot if that would be feasible?

[Aarsh-Wankar]

I have added an extra argument inner_depth to the debugprint function, which controls the depth of the inner graph to be printed. For example:

import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import get_default_mode

n = pt.iscalar("n")
x0 = pt.vector("x0")
xs, _ = pytensor.scan(lambda xtm1: xtm1 + 1, outputs_info=[x0], n_steps=n)

mode = get_default_mode().including("scan_save_mem")
fn = pytensor.function([n, x0], xs, mode=mode, on_unused_input="ignore")
fn.dprint(inner_depth = 1)

# Subtensor{start:stop} [id A] 9
#  ├─ Scan{scan_fn, while_loop=False, inplace=all} [id B] 8
#  │  ├─ n [id C]
#  │  └─ SetSubtensor{:stop} [id D] 7
#  │     ├─ AllocEmpty{dtype='float64'} [id E] 6
#  │     │  ├─ Composite{...}.2 [id F] 0
#  │     │  │  └─ n [id C]
#  │     │  └─ Shape_i{0} [id G] 5
#  │     │     └─ x0 [id H]
#  │     ├─ Unbroadcast{0} [id I] 4
#  │     │  └─ ExpandDims{axis=0} [id J] 3
#  │     │     └─ x0 [id H]
#  │     └─ 1 [id K]
#  ├─ ScalarFromTensor [id L] 2
#  │  └─ Composite{...}.1 [id F] 0
#  │     └─ ···
#  └─ ScalarFromTensor [id M] 1
#     └─ Composite{...}.0 [id F] 0
#        └─ ···

# Inner graphs:

# Scan{scan_fn, while_loop=False, inplace=all} [id B]
#  ← Add [id N]

# Composite{...} [id F]
#  ← add [id O] 'o0'
#  ← add [id P] 'o1'
#  ← add [id Q] 'o2'

fn.dprint(inner_depth = 2)

# ...
# Inner graphs:

# Scan{scan_fn, while_loop=False, inplace=all} [id B]
#  ← Add [id N]
#     ├─ [1.] [id O]
#     └─ *0-<Vector(float64, shape=(?,))> [id P] -> [id D]

# Composite{...} [id F]
#  ← add [id Q] 'o0'
#     ├─ sub [id R]
#     └─ maximum [id S] 't13'
#  ← add [id T] 'o1'
#     ├─ sub [id U]
#     └─ maximum [id S] 't13'
#        └─ ···
#  ← add [id V] 'o2'
#     ├─ Switch [id W]
#     └─ 1 [id X]

Are these changes fine?

[ricardoV94]

depth and inner_depth are not the same. Depth still makes sense alongside inner_depth. Inner depth tells how many inner graphs to step into, depth tells how many ops in a graph (or inner garph) to step into