Implement rich debugprint format via file="rich"

[williambdean]

Closes #1034

Adds file="rich" as a new sentinel to debugprint, returning a
rich.tree.Tree instead of printing text.

import pytensor.tensor as pt
from pytensor.printing import debugprint
from rich.console import Console

x = pt.vector("x")
tree = debugprint([x.mean(), x.std()], file="rich")

console = Console()
console.print(tree)

True_div [id A] 'mean'
├── Sum{axes=None} [id B]
│   └── x [id C]
└── Subtensor{i} [id D]
    ├── Cast{float64} [id E]
    │   └── Shape [id F]
    │       └── x [id C]
    └── 0 [id G]
Sqrt [id H] 'std'
└── True_div [id I] 'var'
    ├── Sum{axis=0} [id J]
    │   └── Pow [id K]
    │       ├── Sub [id L]
    │       │   ├── x [id C]
    │       │   └── True_div [id M] 'mean'
    │       │       ├── ExpandDims{axis=0} [id N]
    │       │       │   └── Sum{axis=0} [id O]
    │       │       │       └── x [id C]
    │       │       └── ExpandDims{axis=0} [id P]
    │       │           └── Subtensor{i} [id Q]
    │       │               ├── Cast{float64} [id R]
    │       │               │   └── Shape [id S]
    │       │               │       └── x [id C]
    │       │               └── 0 [id T]
    │       └── ExpandDims{axis=0} [id U]
    │           └── 2.0 [id V]
    └── Subtensor{i} [id W]
        ├── Cast{float64} [id X]
        │   └── Shape [id Y]
        │       └── Pow [id K] ···
        └── 0 [id Z]

rich is added as a hard dependency. The refactor also extracts
GraphNode, _assign_id, _build_label, and _iter_graph_nodes from
the monolithic _debugprint function, which both the text and rich
renderers now consume.

The "rich" sentinel slots in alongside "str" and None — the text
output path is unchanged.

[williambdean]

The GraphNode / _iter_graph_nodes refactor also makes it straightforward to build other graph renderers on top of the same traversal. As a sketch, here is a Mermaid renderer using the same (x * 2).sum() graph from the tests:

Mermaid renderer sketch

import re
import pytensor.tensor as pt
from pytensor.printing import _build_label, _iter_graph_nodes

def build_mermaid(var) -> str:
    done: dict = {}
    used_ids: dict = {}
    node_defs: dict[str, str] = {}
    edges: list[tuple[str, str]] = []
    apply_to_mid: dict[int, str] = {}

    for gnode in _iter_graph_nodes(var, done=done):
        label = _build_label(
            gnode, done=done, used_ids=used_ids, id_type="CHAR",
            print_type=False, print_shape=False, print_destroy_map=False,
            print_view_map=False, print_op_info=False, op_information={},
        )
        mid = re.search(r"\[id ([^\]]+)\]", label).group(1)
        if mid not in node_defs:
            node_defs[mid] = label
        if gnode.parent_node is not None:
            parent_mid = apply_to_mid.get(id(gnode.parent_node))
            if parent_mid is not None:
                edges.append((parent_mid, mid))
        if gnode.var.owner is not None:
            apply_to_mid[id(gnode.var.owner)] = mid

    lines = ["graph TD"]
    for mid, label in node_defs.items():
        lines.append(f'    {mid}["{label.replace(chr(34), "#quot;")}"]')
    for src, dst in edges:
        lines.append(f"    {src} --> {dst}")
    return "\n".join(lines)

x = pt.dvector("x")
print(build_mermaid((x * 2).sum()))

Output:

graph TD
    A["Sum{axes=None} [id A]"]
    B["Mul [id B]"]
    C["x [id C]"]
    D["ExpandDims{axis=0} [id D]"]
    E["2 [id E]"]
    A --> B
    B --> C
    B --> D
    D --> E

Loading

This would also address #1488.

[williambdean]

Another one from the mermaid POC

x = pt.dvector("x")
y = (x * 2).sum()

z = pt.stack([y.mean(), y.std()])

graph TD
    A["MakeVector{dtype='float64'} [id A]"]
    B["True_div [id B] 'mean'"]
    C["Sum{axes=None} [id C]"]
    D["Sum{axes=None} [id D]"]
    E["Mul [id E]"]
    F["x [id F]"]
    G["ExpandDims{axis=0} [id G]"]
    H["2 [id H]"]
    I["Cast{float64} [id I]"]
    J["1 [id J]"]
    K["Sqrt [id K] 'std'"]
    L["True_div [id L] 'var'"]
    M["Sum{axes=[]} [id M]"]
    N["Pow [id N]"]
    O["Sub [id O]"]
    P["True_div [id P] 'mean'"]
    Q["Sum{axes=[]} [id Q]"]
    R["Cast{float64} [id R]"]
    S["1 [id S]"]
    T["2.0 [id T]"]
    U["Cast{float64} [id U]"]
    V["1 [id V]"]
    A --> B
    B --> C
    C --> D
    D --> E
    E --> F
    E --> G
    G --> H
    B --> I
    I --> J
    A --> K
    K --> L
    L --> M
    M --> N
    N --> O
    O --> D
    D --> D
    O --> P
    P --> Q
    Q --> D
    D --> D
    P --> R
    R --> S
    N --> T
    L --> U
    U --> V

Loading

[ricardoV94]

Can I see how the rich ones look like :D ?

[williambdean]

from rich import print as rprint
x = pt.dvector("x")
y = (x * 2).sum()

z = pt.stack([y.mean(), y.std()])

tree = debugprint(z, file="rich")
rprint(tree)

Result:

MakeVector{dtype='float64'} [id A]
├── True_div [id B] 'mean'
│   ├── Sum{axes=None} [id C]
│   │   └── Sum{axes=None} [id D]
│   │       └── Mul [id E]
│   │           ├── x [id F]
│   │           └── ExpandDims{axis=0} [id G]
│   │               └── 2 [id H]
│   └── Cast{float64} [id I]
│       └── 1 [id J]
└── Sqrt [id K] 'std'
    └── True_div [id L] 'var'
        ├── Sum{axes=[]} [id M]
        │   └── Pow [id N]
        │       ├── Sub [id O]
        │       │   ├── Sum{axes=None} [id D]
        │       │   │   └── Sum{axes=None} [id D] ···
        │       │   └── True_div [id P] 'mean'
        │       │       ├── Sum{axes=[]} [id Q]
        │       │       │   └── Sum{axes=None} [id D]
        │       │       │       └── Sum{axes=None} [id D] ···
        │       │       └── Cast{float64} [id R]
        │       │           └── 1 [id S]
        │       └── 2.0 [id T]
        └── Cast{float64} [id U]
            └── 1 [id V]

[ricardoV94]

That's just text to me, how is it different than the default?

[ricardoV94]

One difference it shows the nested one grayed out... but it's duplicated?

        │       └── Pow [id D]
        │           └── Pow [id D] ···

[ricardoV94]

Would be nice to do some more styling (now that we can). Like coloring these repeated nodes to more easly see where they come from. Pow [id D] would have a unique color

[ricardoV94]

Also no inner graphs :( ?

In [10]: pytensor.OpFromGraph([x], [x.std()])(x).dprint(file="rich")
Out[10]: 
OpFromGraph{inline=False} [id A]
└── x [id B]

[williambdean]

How about something like:

already seen

inner graph

[ricardoV94]

That's good but is the duplication unavoidable? Should be just seen node then ...

[williambdean]

Can be this:

[williambdean]

Pushing up. Cycles through list of colors. Personally find the red looks like an error. Thoughts / preferences?

Reference
https://rich.readthedocs.io/en/stable/appendix/colors.html#standard-colors

[ricardoV94]

I like it, feel free to skip red, but it doesn't bother me.

I'm a bit OCD on the ellipsis not being it's own line, but may be just habit.

OTOH we have a lot of attributes we write inline like destroy map, names, op_info, type, shape that can become long and hide the ellipsis or make it harder to spot

[williambdean]

[williambdean]

OTOH we have a lot of attributes we write inline like destroy map, names, op_info, type, shape that can become long and hide the ellipsis or make it harder to spot

Was this some additional feature? Something to follow up on?

[ricardoV94]

Was this some additional feature?

This already exists in dprint: dprint(print_type=True, print_memory_map=True, print_op_info=True). It's fine if the rich version doesn't currently handle them, but my point was I think I want ... to show below the repeated node, not after it, because we already append (or will append) a lot of information next to it and it's hard to note the ellipsis at the end?

Does it make sense? If making the ellipsis an item on next line is too hard, let's just go with what we have now.

[ricardoV94]

One last question, since we have so much color control, can we render the inner graphs a bit lighter (lower contrast?). To make it easier to scroll into the original graph

[williambdean]

Something like this check out to expectation?

[williambdean]

or completely dimmed?

[ricardoV94]

Just the |- ... (not repeating the node name and id and all that) (like file="str" does). The colors were nice, don't go back on me there

The dimmed inner graph have to think maybe let's not do that for now. (use same style as outer)

[williambdean]

here is the latest

[ricardoV94]

gorgeous.

Does ellipsis color have any meaning? I would leave then gray (or same color as parent) if not

Let's not make the inner graphs less contrast, use exactly the same styling /rules as outer graph. It's something I went back on after first asking.

[williambdean]

Does ellipsis color have any meaning? I would leave then gray (or same color as parent) if not

ellipsis color is same as the parent. Would also be dimmed if part of the InnerGraph

Let's not make the inner graphs less contrast, use exactly the same styling /rules as outer graph. It's something I went back on after first asking.

So not dimmed anymore? How about just bolding the Inner Graph header?

[ricardoV94]

How about just bolding the Inner Graph header?

Yeah that sounds better.

ellipsis color is same as the parent. Would also be dimmed if part of the InnerGraph

That's not what it shows in your last plot?

You also have double ellipsis under. I can guess what's happening but not sure it's good. The simplest (visual wise) is to have ellipsis be regular text color (no color), and only a single ellpsis under a repeated node (regardless of how many inputs it has.

[williambdean]

So no to :

└── Sub [id L]
    ├── ···
    └── ···

and consolidate to:

└── Sub [id L]
    └── ···

Is that your preference?

[ricardoV94]

Yes! That's the same thing regular dprint does

[williambdean]

Looks like this now:

However, don't think the colors add any value now ... They mean nothing

[williambdean]

Don't think this one seems right. i.e. no children but reused?

[williambdean]

file="str" wouldn't have the ... so that needs to change.

However, should color still be used in this case?

[ricardoV94]

Don't think this one seems right. i.e. no children but reused?

Yeah, if it's a root variable don't need to use ellipsis below, although the color is still nice imo to see the root variable is used in multiple places (root variables only used in one place can stay regular gray).

Let me know if this complicates things a lot, if yes don't bother with the color for this case

[williambdean]

Fixed with 76d5878

[ricardoV94]

Tested it locally, it's awesome. I think the only thing not supported is print_fgraph_inputs, but whatever. Very good as is

[ricardoV94]

@williambdean ready to merge?

[ricardoV94]

Got this bot review if you care to look (I know you could get it for yourself):

PR Review

Overview

The PR does two things:

1. Refactors the recursive _debugprint into a _iter_graph_nodes generator that yields GraphNode records, with _build_label and _assign_id extracted as pure helpers.
2. Adds _build_rich_tree which consumes the same generator to build a rich.tree.Tree. Repeated nodes across a DAG get a matching color, and their second occurrence is replaced by a ··· sentinel under a colored label.

The text renderer now delegates to the generator, so the refactor is exercised by every pre-existing debugprint test.

Strengths

• Shared traversal logic. Renderers consume a stream of GraphNode records — no duplicated graph-walking. This is the right direction.
• Lazy rich import. _build_rich_tree imports inside the function with a clear error; rich stays a soft dep (despite the commit message "Add rich as a hard dependency," it isn't in pyproject.toml dependencies — which is correct, but worth aligning with an extras entry e.g. pytensor[rich]).
• Test structure. 17 tests exercise structural contracts rather than rendered strings — keeps rich's rendering out of the test surface. Markup-escaping test (test_markup_escaping) and shared-across-outputs test (test_shared_node_colored_across_outputs) are particularly good.
• Markup escaping. rich.markup.escape(label) is applied before decorating — correct order.

Issues / Suggestions

Duplication in _build_rich_tree (printing.py:1127–1341)

Three nearly identical loops consume _iter_graph_nodes — top-level outputs, inner-graph headers, and inner-graph outputs. Each re-implements the same is_already_done / is_repeat / _build_label / node_key / depth-stack logic (~60 lines × 3). Extracting a helper like _emit_tree_from_iter(root_tree, var, **iter_kwargs) would cut ~150 lines and eliminate the risk of these paths drifting.

GraphNode reconstructed, not updated

In _iter_graph_nodes the GraphNode is built once, then replaced wholesale in the leaf-repeat and apply-already-done branches (printing.py:380, 420). dataclasses.replace(gnode, is_already_done=True) would be clearer, or just mutate.

Parameter explosion

_iter_graph_nodes takes 14 parameters; _build_label takes 10. A small PrintContext dataclass carrying (id_type, print_type, print_shape, print_destroy_map, print_view_map, print_op_info, op_information) would clean up both signatures and the many _build_label(gnode, done=..., used_ids=..., id_type=..., ...) call sites.

Missing profile output in rich renderer

The text renderer still appends --> 1.23e-04s 5.0% ... when profile.apply_time is present (printing.py:712–725), but _build_rich_tree drops it. If this is intentional, worth a comment; otherwise it's a behavioral gap vs the text renderer.

done tracking changed for leaves

_iter_graph_nodes now writes leaves into done (printing.py:378) to detect leaf repeats for coloring. Original _debugprint only wrote via get_id_str. Behavior appears preserved for the text path (the new is_already_done flag is ignored there), but this is a subtle semantic change — please double-check that external callers passing in a pre-populated done dict still work.

Type-annotation nits

• dict | None → dict[Apply, str] | None / dict[Variable, Variable] | None at printing.py:302, 306, 311.
• -> "str | TextIO | Any": # rich.tree.Tree when file="rich" — consider from typing import TYPE_CHECKING + a real annotation instead of the stringified Any.

Repeated node_key idiom

node_key = gnode.var.owner if gnode.var.owner is not None else gnode.var appears 6+ times. Make it a GraphNode.key property.

Risk

• Low risk for existing text-debugprint callers: refactor is behavior-preserving and covered by existing tests.
• Rich path is new, opt-in, and soft-imported — no impact on users who don't pass file="rich".
• Main concern is maintainability: the three copy-pasted loops in _build_rich_tree will invite bugs the next time someone touches rich rendering.

[williambdean]

ready to merge 🚀

[ricardoV94]

Thanks @williambdean