ENH: Expose a 'stop at named variables' mode in `pprint` and add LaTeX printers

ENH: Expose a "stop at named variables" mode in pprint and add LaTeX printers

Repo: pymc-devs/pytensor
Related: pymc-devs/pymc issue "[ENH: Render full symbolic expressions for Deterministics ...]" (filed first — it consumes this feature): https://github.com/pymc-devs/pymc/issues/8407

Motivation

PyTensor's pprint is the go-to tool for showing a symbolic expression as math-like text, but it has one blind spot: it never stops at named intermediate variables. Given

import pytensor.tensor as pt
from pytensor.printing import pprint

x = pt.vector("x")
beta = pt.vector("beta")
sigma = pt.vector("sigma")
mu = x * beta + pt.log(sigma**2) + 3
print(pprint(mu))

you get the fully expanded form with every leaf spelled out — fine for debugging internals, but useless when x, beta, and sigma are meaningful names that should be treated as atoms:

(((x * beta) + Log((sigma ** 2))) + 3)   # what we want
# vs. full RV/op expansion when leaves are e.g. RandomVariables

This matters whenever a graph's named variables correspond to user-facing concepts — which is exactly the situation in downstream libraries like PyMC, where model variables have names and deterministics should render as mu = x * beta + log(sigma**2) + 3 rather than f(...). A concrete proposal to use this for richer PyMC model representations (plain text and LaTeX) is filed over in pymc (pymc-devs/pymc#8407); it currently works around this gap by cloning the global printer from library code.

Separately: there is no LaTeX output anywhere in PyTensor. The printer infrastructure (OperatorPrinter, PatternPrinter, FunctionPrinter) is format-agnostic — patterns like \frac{%(0)s}{%(1)s} work today — so an official LaTeX dialect is a natural extension rather than new machinery. This would benefit anyone exporting symbolic math to docs, notebooks, or papers.

Requested features

1. Named-variable leaf mode for pprint

The logic already exists internally: PPrinter.process_graph (printing.py, ~line 1840) clones itself with

pprinter.clone_assign(lambda pstate, r: r.name is not None and r is not current, leaf_printer)

but this behavior is not reachable via a public API on pprint.process. A small addition such as

pprint(expr, named_leaves=True)          # stop at any variable with a name
pprint(expr, named_leaves={"beta", "sigma"})  # stop at specific names/vars

would expose it. One implementation note discovered while prototyping from outside: entries in PPrinter.printers_dict (e.g. DimShuffle) take precedence over condition-based assignments, so an override API should account for that ordering.

2. LaTeX printers

A latex_pprint (or a dialect="latex" option) registering:

• arithmetic: +, - (unary/binary), implicit or \cdot multiplication, \frac{a}{b}, {a}^{b}
• common functions: \exp, \log, \sqrt, trig, etc.
• reductions (\sum), matrix products (\cdot / \times)
• subtensor/indexing: plain text already renders X[:, 0] nicely via registered Subtensor printers; LaTeX needs an equivalent (e.g. X_{[:,\,0]}) — currently absent, so indexing degrades to an opaque fallback
• graceful generic fallback: \operatorname{name}(args) for unregistered ops, so nothing ever crashes

The existing precedence-aware parenthesization comes free by reusing OperatorPrinter/PatternPrinter; only leaf naming/escaping policy needs deciding (e.g. greek normalization beta → \beta).

Why upstream rather than library-side clones

Both behaviors are achievable today via public APIs (clone_assign), but each downstream consumer re-implementing printer clones leads to divergent output conventions and duplicated maintenance. Putting them in pytensor.printing gives one canonical implementation next to the existing printer classes, where op coverage can grow incrementally.

Start in printing.py at PPrinter.process_graph and the public pprint.process API, then review OperatorPrinter, PatternPrinter, and FunctionPrinter. Implement a public named-leaf mode while accounting for printers_dict precedence, and add a LaTeX dialect with arithmetic, functions, reductions, indexing, and generic fallback printers. Done means both requested output modes work without downstream printer clones.