diff --git a/pymc/printing.py b/pymc/printing.py index 01988200dc..7d53168d60 100644 --- a/pymc/printing.py +++ b/pymc/printing.py @@ -16,18 +16,37 @@ import re import sys +from collections import deque from collections.abc import Iterable from functools import partial +from numbers import Integral import numpy as np import pytensor.tensor as pt from pytensor.compile import SharedVariable +from pytensor.compile.builders import OpFromGraph +from pytensor.compile.ops import ViewOp from pytensor.graph.basic import Constant, Variable -from pytensor.graph.traversal import walk +from pytensor.graph.traversal import ancestors, walk from pytensor.graph.type import HasShape -from pytensor.tensor.elemwise import DimShuffle +from pytensor.printing import ( + FunctionPrinter, + OperatorPrinter, + PatternPrinter, + PPrinter, + Printer, + PrinterState, + set_precedence, +) +from pytensor.printing import pprint as _pytensor_pprint +from pytensor.scan.op import Scan +from pytensor.tensor.blockwise import Blockwise +from pytensor.tensor.elemwise import DimShuffle, Elemwise +from pytensor.tensor.math import Dot, Sum +from pytensor.tensor.random.op import RNGConsumerOp from pytensor.tensor.random.type import RandomType +from pytensor.tensor.subtensor import AdvancedSubtensor, Subtensor from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable from rich.box import SIMPLE_HEAD @@ -139,11 +158,36 @@ def str_for_data_var( return rf"{print_name} = Data" -def str_for_model(model: Model, formatting: str = "plain", include_params: bool = True) -> str: +def str_for_model( + model: Model, + formatting: str = "plain", + include_params: bool = True, + deterministic_exprs: bool = False, +) -> str: """Make a human-readable string representation of Model. This lists all random variables and their distributions, optionally including parameter values. + + Parameters + ---------- + model + The model to represent. + formatting + Either "plain" or "latex". + include_params + Whether to include parameter values. + deterministic_exprs + If True, Deterministics and Potentials render their full symbolic + expression instead of an opaque ``f(inputs)`` placeholder. Traversal + stops at named model variables, matched by identity, which are + rendered by name. Purely a graph traversal: no rewrites, compilation, + or evaluation happen. ``include_params=False`` takes precedence, + falling back to the opaque ``Deterministic``/``Potential`` rendering. + Expressions are bounded: unnamed subgraphs deeper than 64 levels + render as the opaque placeholder, and if more than 1000 nodes would + still be rendered, the whole expression degrades to it. Naming + intermediates keeps large models fully expanded. """ named_vars: set[Variable] = set() named_vars.update(model.data_vars) @@ -161,6 +205,7 @@ def str_for_model(model: Model, formatting: str = "plain", include_params: bool formatting=formatting, include_params=include_params, named_vars=named_vars, + deterministic_exprs=deterministic_exprs, ) sfdv = partial(str_for_data_var, formatting=formatting, include_params=include_params) @@ -213,18 +258,35 @@ def str_for_potential_or_deterministic( include_params: bool = True, dist_name: str = "Deterministic", named_vars: set[Variable] | None = None, + deterministic_exprs: bool = False, ) -> str: """Make a human-readable string representation of a Deterministic or Potential in a model. This can be either LaTeX or plain, optionally with distribution parameter values included. + + ``deterministic_exprs`` only takes effect when ``include_params=True``; + otherwise the opaque ``Deterministic``/``Potential`` form is rendered. + If ``named_vars`` is omitted, every named variable reachable from ``var`` + is treated as a known variable and rendered by name, so standalone calls + never inline distributions into the expression. + + Expression bodies are bounded: unnamed subgraphs deeper than 64 levels + render as the opaque placeholder, and if more than 1000 nodes would + still be rendered, the whole expression degrades to it. """ if named_vars is None: - named_vars = set() + named_vars = {v for v in ancestors([var]) if v.name is not None} print_name = var.name if var.name is not None else "" sep_plain = "~" if dist_name == "Potential" else "=" sep_latex = r"\sim" if dist_name == "Potential" else "=" + if deterministic_exprs and include_params: + expr = _str_for_expression_body(var, formatting=formatting, named_vars=named_vars) + if "latex" in formatting: + latex_name = r"\text{" + _latex_escape(print_name.strip("$")) + "}" + return rf"${latex_name} {sep_latex} {expr}$" + return rf"{print_name} {sep_plain} {expr}" if "latex" in formatting: print_name = r"\text{" + _latex_escape(print_name.strip("$")) + "}" if include_params: @@ -325,6 +387,587 @@ def _expand(x): return r"f(" + ", ".join([n.strip("$") for n in names]) + ")" +class _TransparentFirstInputPrinter(Printer): + """Render a unary ``ViewOp`` as its first input. + + Identity wrappers around deterministics and other unnamed view ops carry + no mathematical content, so they are skipped during rendering. + """ + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + r = pstate.pprinter.process(output.owner.inputs[0], pstate) + pstate.memo[output] = r + return r + + +def _dimshuffle_is_broadcast_only(op) -> bool: + """True if a DimShuffle only inserts 'x' axes or drops broadcastable ones. + + Surviving axes must stay in their original order, i.e. nothing is + permuted and dropping the node cannot change the printed expression. + """ + last = -1 + for o in op.new_order: + if o == "x": + continue + if not isinstance(o, Integral) or o <= last: + return False + last = int(o) + return True + + +class _DimShufflePrinter(Printer): + """Render DimShuffle without hiding real axis permutations. + + Pure broadcasting renders transparently; a full axis reversal renders as + a transpose (``X.T`` / ``{X}^{T}``); any other permutation shows its axis + order so the printed expression stays faithful to the graph. + """ + + def __init__(self, formatting: str): + self.formatting = formatting + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + op = output.owner.op + with set_precedence(pstate): + arg = pstate.pprinter.process(output.owner.inputs[0], pstate) + if _dimshuffle_is_broadcast_only(op): + r = arg + elif op.new_order == tuple(range(output.type.ndim))[::-1]: + r = rf"{{{arg}}}^{{T}}" if "latex" in self.formatting else rf"{arg}.T" + else: + order = tuple(int(o) for o in op.new_order) + if "latex" in self.formatting: + shown = r",\ ".join(str(o) for o in order) + r = rf"\operatorname{{transpose}}\left({arg},~\text{{order}}=\left({shown}\right)\right)" + else: + r = rf"transpose({arg}, order={order})" + pstate.memo[output] = r + return r + + +def _is_matmul(r) -> bool: + """Matrix products: a plain Dot or a Blockwise wrapping a core Dot.""" + op = getattr(r.owner, "op", None) + if op is None: + return False + return isinstance(op, Dot) or (isinstance(op, Blockwise) and isinstance(op.core_op, Dot)) + + +class _BodyLeafPrinter(Printer): + """Render named model variables by name, without expanding their graphs.""" + + def __init__(self, formatting: str): + self.formatting = formatting + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + # The leaf condition may match an unnamed ViewOp wrapper whose + # unwrapped target is the named variable being rendered. + name = _unwrap_viewops(output).name.strip("$") + if "latex" in self.formatting: + r = rf"\text{{{_latex_escape(name)}}}" + else: + r = name + pstate.memo[output] = r + return r + + +class _BodyDistPrinter(Printer): + """Render anonymous distributions inside bodies as distribution calls. + + Delegates to ``str_for_dist`` so an inline prior such as + ``pm.Normal.dist(0, 1)`` looks like its named counterpart on RV lines, + instead of leaking graph internals (including a nondeterministic RNG + memory address) into the repr. + """ + + def __init__(self, formatting: str, named_vars: set[Variable]): + self.formatting = formatting + self.named_vars = named_vars + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + r = str_for_dist(output, formatting=self.formatting, named_vars=self.named_vars) + # str_for_dist wraps latex in $...$ for standalone use; strip for inline bodies + r = r.strip("$") + pstate.memo[output] = r + return r + + +class _LatexFunctionPrinter(Printer): + r"""Fallback LaTeX rendering: \operatorname{name}(args). + + Distinguishing op parameters are carried into the output so distinct ops + never collapse onto the same rendering: ``Blockwise`` and ``Elemwise`` + are unwrapped to their inner op, and casts show their target dtype. + """ + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + op = output.owner.op + dtype = None + if isinstance(op, Blockwise): + name = type(op.core_op).__name__ + elif isinstance(op, Elemwise): + scalar_op = op.scalar_op + name = type(scalar_op).__name__.lower() + o_type = getattr(scalar_op, "o_type", None) + if o_type is not None: + dtype = o_type.dtype + else: + name = getattr(op, "name", None) or type(op).__name__ + name = re.sub(r"\W+", "", str(name)) or "op" + with set_precedence(pstate): + args = [pstate.pprinter.process(i, pstate) for i in output.owner.inputs] + if dtype is not None: + args.append(rf"\text{{{dtype}}}") + r = rf"\operatorname{{{name}}}\left({r',\ '.join(args)}\right)" + pstate.memo[output] = r + return r + + +class _LatexSumPrinter(Printer): + r"""LaTeX Sum rendering that keeps the reduction axes visible.""" + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + op = output.owner.op + with set_precedence(pstate): + arg = pstate.pprinter.process(output.owner.inputs[0], pstate) + if op.axis is None: + r = rf"\sum\left({arg}\right)" + else: + axes = r",\ ".join(str(int(a)) for a in op.axis) + r = rf"\sum_{{{axes}}}\left({arg}\right)" + pstate.memo[output] = r + return r + + +class _BodyConstantPrinter(Printer): + """Render constants inside expression bodies. + + Scalars and single-element arrays render their value, short 1-D arrays + inline their values, and larger arrays describe their dtype and shape + instead: geometry is the information a reader needs, and the values + would flood the line. This only affects opt-in expression bodies; + ``_str_for_constant_value`` keeps default ``str_for_model`` output + unchanged. + """ + + def __init__(self, formatting: str): + self.formatting = formatting + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + data = output.data + if isinstance(data, np.ndarray): + r = self._render_array(data) + else: + # e.g. NoneConst placeholders inside RV signatures + r = str(data) + pstate.memo[output] = r + return r + + def _render_array(self, data: np.ndarray) -> str: + latex = "latex" in self.formatting + if data.ndim == 0 or data.size == 1: + return _str_for_constant_value(data, self.formatting) + if data.ndim == 1 and data.size <= _MAX_INLINE_CONST_ELEMENTS: + sep = r",\ " if latex else ", " + values = sep.join(f"{v:.3g}" for v in data) + return rf"\left[{values}\right]" if latex else f"[{values}]" + shape = ", ".join(map(str, data.shape)) + if not latex: + return f"" + number_set = {"f": r"\mathbb{R}", "i": r"\mathbb{Z}", "u": r"\mathbb{Z}"}.get( + data.dtype.kind + ) + if number_set is None: + return rf"\text{{" + dims = r" \times ".join(map(str, data.shape)) + return rf"\text{{}} \in {number_set}^{{{dims}}}" + + +class _OwnerlessLeafPrinter(Printer): + r"""Render ownerless leaves consistently across formats. + + A named-but-not-in-model variable (e.g. an external shared container) + renders by name; an anonymous one by its type, mirroring plain text. + """ + + def __init__(self, formatting: str): + self.formatting = formatting + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + name = getattr(output, "name", None) + r = name.strip("$") if name is not None else f"<{output.type}>" + if "latex" in self.formatting: + r = rf"\text{{{_latex_escape(r)}}}" + pstate.memo[output] = r + return r + + +def _is_unary_viewop(r) -> bool: + return r.owner is not None and isinstance(r.owner.op, ViewOp) and len(r.owner.inputs) == 1 + + +def _is_random_op(r) -> bool: + # Covers both pytensor RandomVariables and PyMC SymbolicRandomVariables, + # which also derive from RNGConsumerOp. + return r.owner is not None and isinstance(r.owner.op, RNGConsumerOp) + + +def _hides_inner_graph(r) -> bool: + """True for nodes whose rendering would leak inner-graph machinery. + + Ops like ``Scan`` and ``OpFromGraph`` carry their own subgraph: printing + their inputs shows allocation machinery while never showing the loop + body. Slicing such a result also leaks the slice's index machinery, so it + is hidden as well. + """ + if r.owner is None: + return False + op = r.owner.op + if isinstance(op, Scan | OpFromGraph): + return True + return isinstance(op, Subtensor | AdvancedSubtensor) and _hides_inner_graph(r.owner.inputs[0]) + + +class _InnerGraphOpPrinter(Printer): + """Render inner-graph ops as an opaque ``f(...)`` placeholder.""" + + def __init__(self, formatting: str, named_vars: set[Variable]): + self.formatting = formatting + self.named_vars = named_vars + + def process(self, output, pstate): + if output in pstate.memo: + return pstate.memo[output] + r = _str_for_expression(output, self.formatting, self.named_vars) + pstate.memo[output] = r + return r + + +class _LeafGuardPrinter(Printer): + """Intercept a dict-keyed printer so named model leaves still win. + + ``PPrinter.process`` checks dict-keyed registrations before every + condition rule, so an op instance shared between a named leaf's owner and + ordinary nodes would otherwise expand past the leaf boundary. + """ + + def __init__(self, inner: Printer, named_leaf_condition, leaf_printer: Printer): + self.inner = inner + self.named_leaf_condition = named_leaf_condition + self.leaf_printer = leaf_printer + + def process(self, output, pstate): + if self.named_leaf_condition(pstate, output): + return self.leaf_printer.process(output, pstate) + return self.inner.process(output, pstate) + + +def _shield_named_leaves( + printer: PPrinter, named_leaf_condition, leaf_printer: Printer, named_vars: set[Variable] +): + """Wrap dict-keyed registrations that could swallow a named leaf.""" + ops = {v.owner.op for v in named_vars if getattr(v, "owner", None) is not None} + for op in ops: + for key in (op, type(op)): + inner = printer.printers_dict.get(key) + if inner is not None: + printer.printers_dict[key] = _LeafGuardPrinter( + inner, named_leaf_condition, leaf_printer + ) + break + + +def _unwrap_viewops(r): + """Descend through identity wrappers, but never past a named variable. + + Named Deterministics are themselves unary ``ViewOp`` outputs + (``view_op(var, name=...)``), so unwrapping must treat any named node as + a hard boundary. + """ + while getattr(r, "name", None) is None and _is_unary_viewop(r): + r = r.owner.inputs[0] + return r + + +def _make_plain_body_printer(named_leaf_condition, named_vars: set[Variable]) -> PPrinter: + """Clone the global pytensor printer with model-aware overrides. + + Inheriting the global printer's registrations gives plain-text coverage of + many ops for free. + + Anonymous distributions are rendered via ``str_for_dist`` (like named RV + lines) rather than the inherited raw-RNG rendering, which leaks a + nondeterministic memory address. Matrix products print as ``@`` (the + inherited registration emits a LaTeX escape into plain output), real axis + permutations stay visible, and inner-graph ops degrade to opaque + placeholders. + + Dict-keyed registrations outrank every condition rule in + ``PPrinter.process``, so ``DimShuffle`` and ``Scan`` are overridden in + kind, and entries shared with a named leaf's owner are wrapped via + ``_shield_named_leaves``. + """ + leaf_printer = _BodyLeafPrinter("plain") + printer = _pytensor_pprint.clone_assign(DimShuffle, _DimShufflePrinter("plain")) + printer = printer.clone_assign(Scan, _InnerGraphOpPrinter("plain", named_vars)) + at_printer = OperatorPrinter("@", -1, "left") + # pytensor registers its Dot singleton by instance, which outranks the + # class-keyed entry assigned next; override it in kind + printer = printer.clone_assign(Dot, at_printer) + for k in [k for k in printer.printers_dict if isinstance(k, Dot)]: + printer.printers_dict[k] = at_printer + printer = printer.clone_assign( + lambda pstate, r: _is_matmul(r), + OperatorPrinter("@", -1, "left"), + ) + printer = printer.clone_assign( + lambda pstate, r: _is_random_op(r), + _BodyDistPrinter("plain", named_vars), + ) + printer = printer.clone_assign( + lambda pstate, r: _is_unary_viewop(r), + _TransparentFirstInputPrinter(), + ) + printer = printer.clone_assign( + lambda pstate, r: _hides_inner_graph(r), + _InnerGraphOpPrinter("plain", named_vars), + ) + printer = printer.clone_assign( + lambda pstate, r: isinstance(r, Constant), + _BodyConstantPrinter("plain"), + ) + printer = printer.clone_assign( + lambda pstate, r: r.owner is None and not isinstance(r, Constant), + _OwnerlessLeafPrinter("plain"), + ) + printer = printer.clone_assign(named_leaf_condition, leaf_printer) + _shield_named_leaves(printer, named_leaf_condition, leaf_printer, named_vars) + return printer + + +def _make_latex_body_printer(named_leaf_condition, named_vars: set[Variable]) -> PPrinter: + r"""A fresh printer rendering expression bodies as LaTeX. + + Priority is the reverse of assignment order (``assign`` inserts at head). + Ops without a dedicated registration degrade gracefully to + ``\\operatorname{name}(args)``. + + Dict-keyed registrations outrank every condition rule in + ``PPrinter.process``. Condition-based ViewOp handling must therefore stay + below the named-leaf rule, which resolves ViewOp wrappers itself; + DimShuffle is overridden with a dict key because pytensor's own dict + entry would otherwise take precedence (DimShuffle outputs are never + named, so it cannot preempt the leaf rule). + """ + leaf_printer = _BodyLeafPrinter("latex") + printer = PPrinter() + printer.assign(lambda pstate, r: True, _LatexFunctionPrinter()) # lowest priority + printer.assign(lambda pstate, r: r.owner.op is pt.exp, FunctionPrinter([r"\exp"])) + printer.assign(lambda pstate, r: r.owner.op is pt.log, FunctionPrinter([r"\log"])) + printer.assign( + lambda pstate, r: r.owner.op is pt.sqrt, + PatternPrinter((r"\sqrt{%(0)s}",)), + ) + printer.assign(lambda pstate, r: r.owner.op is pt.sin, FunctionPrinter([r"\sin"])) + printer.assign(lambda pstate, r: r.owner.op is pt.cos, FunctionPrinter([r"\cos"])) + printer.assign(lambda pstate, r: r.owner.op is pt.tanh, FunctionPrinter([r"\tanh"])) + printer.assign( + lambda pstate, r: isinstance(r.owner.op, Sum), + _LatexSumPrinter(), + ) + printer.assign( + lambda pstate, r: _is_matmul(r), + PatternPrinter((r"(%(0)s \cdot %(1)s)",)), + ) + printer.assign( + lambda pstate, r: r.owner.op is pt.true_div, + PatternPrinter((r"\frac{%(0)s}{%(1)s}",)), + ) + printer.assign( + lambda pstate, r: r.owner.op is pt.pow, + PatternPrinter((r"{%(0)s}^{%(1)s}",)), + ) + printer.assign(lambda pstate, r: r.owner.op is pt.neg, OperatorPrinter("-", 0, "either")) + printer.assign(lambda pstate, r: r.owner.op is pt.sub, OperatorPrinter("-", -2, "left")) + printer.assign( + lambda pstate, r: r.owner.op is pt.add, + OperatorPrinter("+", -2, "either"), + ) + printer.assign( + lambda pstate, r: r.owner.op is pt.mul, + OperatorPrinter(r"\cdot", -1, "either"), + ) + # Anonymous distributions render as distribution calls; without this they + # would fall through to conditions that assume r.owner is not None and + # crash on their ownerless rng/size inputs. + # Ownerless leaves (e.g. unnamed shared variables referenced directly) + # must terminate here: the operator conditions below assume r.owner. + printer.assign( + lambda pstate, r: r.owner is None and not isinstance(r, Constant), + _OwnerlessLeafPrinter("latex"), + ) + printer.assign( + lambda pstate, r: _is_random_op(r), + _BodyDistPrinter("latex", named_vars), + ) + # Condition-based ViewOp handling: must stay below the named-leaf rule, + # which resolves ViewOp wrappers itself. + printer.assign(lambda pstate, r: _is_unary_viewop(r), _TransparentFirstInputPrinter()) + printer.assign(DimShuffle, _DimShufflePrinter("latex")) + # Inner-graph ops (and slices thereof) hide their machinery behind the + # opaque placeholder; placed above the generic fallback so Scan results + # sliced through Subtensor conditions collapse entirely. + printer.assign( + lambda pstate, r: _hides_inner_graph(r), + _InnerGraphOpPrinter("latex", named_vars), + ) + printer.assign( + lambda pstate, r: isinstance(r, Constant), + _BodyConstantPrinter("latex"), + ) + printer.assign(named_leaf_condition, leaf_printer) # highest priority + _shield_named_leaves(printer, named_leaf_condition, leaf_printer, named_vars) + return printer + + +def _strip_outer_parens(s: str) -> str: + while s.startswith("(") and s.endswith(")") and _parens_balanced(s[1:-1]): + s = s[1:-1] + return s + + +def _parens_balanced(s: str) -> bool: + depth = 0 + for char in s: + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth < 0: + return False + return depth == 0 + + +_MAX_EXPR_NODES = 1000 +_MAX_EXPR_DEPTH = 64 +_MAX_INLINE_CONST_ELEMENTS = 4 + + +def _is_expr_leaf(r: Variable, leaf_vars: set[Variable]) -> bool: + """True for nodes rendered atomically, without expanding their inputs.""" + return ( + _unwrap_viewops(r) in leaf_vars + or isinstance(r, Constant) + or r.owner is None + or _is_random_op(r) + or _hides_inner_graph(r) + ) + + +def _print_plan(body: Variable, leaf_vars: set[Variable]) -> tuple[set[Variable], bool]: + """Decide how to print an expression body within the verbosity budget. + + Printing renders the graph as a tree: pytensor's print memo avoids + recomputing shared subexpressions, not duplicating them, so output size + grows with the number of root-to-node paths, which is exponential for + shared intermediates (an adstock chain reused twice, ``v = v + v`` + loops). Printer recursion depth also follows graph depth. + + Returns ``(cut_roots, fits)``. Nodes at depth ``_MAX_EXPR_DEPTH + 1`` + are cut: they render as the opaque ``f(inputs)`` placeholder via a + pre-seeded print memo, bounding recursion depth. ``fits`` is False when + even after these cuts more than ``_MAX_EXPR_NODES`` nodes would be + rendered; the caller should then fall back to a single placeholder for + the whole expression. + """ + depth = {body: 0} + visited = {body} + queue = deque([body]) + order = [] + while queue: + r = queue.popleft() + order.append(r) + if _is_expr_leaf(r, leaf_vars): + continue + for i in r.owner.inputs: + if i not in visited: + visited.add(i) + depth[i] = depth[r] + 1 + queue.append(i) + + # Minimal cut: seeding only the shallowest nodes past the depth limit + # makes every deeper node unreachable, so each placeholder walk happens + # once instead of once per cut node. + cut_roots = { + r for r, d in depth.items() if d == _MAX_EXPR_DEPTH + 1 and not _is_expr_leaf(r, leaf_vars) + } + + # Occurrences of each node in the rendered tree, counting only paths + # through visible (uncut) parents. ``order`` is breadth-first, so all + # parents are final before a child is accumulated. + occ = {body: 1} + for r in order: + if depth[r] > _MAX_EXPR_DEPTH or _is_expr_leaf(r, leaf_vars): + continue + for i in r.owner.inputs: + occ[i] = occ.get(i, 0) + occ[r] + occurrences = sum(o for r, o in occ.items() if depth[r] <= _MAX_EXPR_DEPTH) + return cut_roots, occurrences <= _MAX_EXPR_NODES + + +def _str_for_expression_body(var: Variable, formatting: str, named_vars: set[Variable]) -> str: + """Render the full symbolic body of an expression graph as text or LaTeX. + + Traversal stops at any model variable in ``named_vars`` (matched by + identity, so unrelated variables that merely share a name are not + mistaken for model variables); those render by name since their + definitions appear elsewhere in the model representation. The rendered + variable itself is excluded so its own body expands. + """ + body = var + while _is_unary_viewop(body) and (body is var or getattr(body, "name", None) is None): + body = body.owner.inputs[0] + + leaf_vars = {v for v in named_vars if v is not var} + + def _named_leaf_condition(pstate, r) -> bool: + return _unwrap_viewops(r) in leaf_vars + + if "latex" in formatting: + printer = _make_latex_body_printer(_named_leaf_condition, named_vars) + else: + printer = _make_plain_body_printer(_named_leaf_condition, named_vars) + + cut_roots, fits = _print_plan(body, leaf_vars) + if not fits: + return _str_for_expression(var, formatting, named_vars) + + pstate = PrinterState(pprinter=printer) + for r in cut_roots: + # Pre-seeding the memo makes every printer stop here, rendering the + # subtree as the opaque placeholder instead of recursing into it. + pstate.memo[r] = _str_for_expression(r, formatting, named_vars) + s = printer.process(body, pstate) + return _strip_outer_parens(s) + + def _latex_text_format(text: str) -> str: if r"\operatorname{" in text: return text diff --git a/tests/test_printing.py b/tests/test_printing.py index ae8372c578..276302ff01 100644 --- a/tests/test_printing.py +++ b/tests/test_printing.py @@ -15,6 +15,8 @@ import re import numpy as np +import pytensor.tensor as pt +import pytest from pytensor.tensor.random import normal from rich.console import Console @@ -568,3 +570,462 @@ def test_model_table(): beta_subject_penalty = Potential(f(beta_subject)) subject[20] """ assert [s.strip() for s in table_txt.splitlines()] == [s.strip() for s in expected.splitlines()] + + +class TestDeterministicExprs: + @staticmethod + def model() -> Model: + with Model() as model: + x = pm.Data("x", 2.0) + sigma = HalfNormal("sigma", sigma=1) + alpha_a = Normal("alpha_a", 0, 1) + mu = Deterministic("mu", alpha_a * x + pt.log(sigma**2) + 3) + Deterministic("eta", mu / (1 + mu)) + Potential("pot", sigma * 2) + Normal("y", mu=mu, sigma=sigma) + return model + + def test_default_repr_unchanged(self): + """Without the flag, deterministics still render opaque f(...) calls.""" + model_text = self.model().str_repr() + assert "mu = Deterministic(f(x, alpha_a, sigma))" in model_text + assert "eta = Deterministic(f(mu))" in model_text + assert "pot ~ Potential(f(sigma))" in model_text + + def test_plain_expression_bodies(self): + model_text = self.model().str_repr(deterministic_exprs=True) + assert "mu = ((alpha_a * x) + Log((sigma ** 2))) + 3" in model_text + # Nested deterministics stop at named variables + assert "eta = mu / (1 + mu)" in model_text + # Potentials get bodies too + assert "pot ~ sigma * 2" in model_text + + def test_latex_expression_bodies(self): + model_tex = self.model().str_repr(formatting="latex", deterministic_exprs=True) + assert ( + r"\text{mu} &= &((\text{alpha\_a} \cdot \text{x}) + \log({\text{sigma}}^{2})) + 3" + in model_tex + ) + assert r"\text{eta} &= &\frac{\text{mu}}{(1 + \text{mu})}" in model_tex + assert r"\text{pot} &\sim & \text{sigma} \cdot 2" in model_tex + + def test_latex_underscore_escaping_in_bodies(self): + model_tex = self.model().str_repr(formatting="latex", deterministic_exprs=True) + assert "\\_" in model_tex + body_tex = model_tex.replace("\\_", "") + assert "_" not in body_tex.replace(r"\_", "") + + def test_standalone_str_for_potential_or_deterministic(self): + from pymc.printing import str_for_potential_or_deterministic + + model = self.model() + mu = next(v for v in model.deterministics if v.name == "mu") + named_vars = set(model.deterministics) | set(model.free_RVs) | set(model.data_vars) + assert ( + str_for_potential_or_deterministic(mu, named_vars=named_vars, deterministic_exprs=True) + == "mu = ((alpha_a * x) + Log((sigma ** 2))) + 3" + ) + assert ( + str_for_potential_or_deterministic(mu, named_vars=named_vars) + == "mu = Deterministic(f(x, alpha_a, sigma))" + ) + + def test_unnamed_viewop_wrapping_named_leaf(self): + # An unnamed ViewOp wrapper around a named variable must render as + # that variable, not crash on its missing name + from pytensor.compile.ops import view_op + + with Model() as model: + s = HalfNormal("s", 1) + Deterministic("d", view_op(s) * 2 + s) + text = model.str_repr(deterministic_exprs=True) + assert "d = (s * 2) + s" in text + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + assert r"\text{d} &= &(\text{s} \cdot 2) + \text{s}" in tex + + def test_latex_elemwise_fallback_and_memoized_nodes(self): + # Unregistered Elemwise ops degrade to \operatorname, and nodes + # referenced more than once render identically at each occurrence + with Model() as model: + s = HalfNormal("s", 1) + X = pm.Data("X", np.eye(2)) + c = pt.as_tensor_variable(2.0) + sig = pt.sigmoid(s) + XT = X.T + Deterministic("d", sig + sig / c + XT @ XT * c) + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + assert r"\operatorname{sigmoid}" in tex + assert tex.count(r"\operatorname{sigmoid}\left(\text{s}\right)") == 2 + assert r"\frac{\operatorname{sigmoid}\left(\text{s}\right)}{2}" in tex + assert r"({\text{X}}^{T} \cdot {\text{X}}^{T}) \cdot 2" in tex + + def test_transpose_rendered_not_deleted(self): + # Adversarial review bug 1: real axis permutations must stay visible + with Model() as model: + X = pm.Data("X", np.eye(3)) + Deterministic("dT", X.T) + Deterministic("dTX", X.T @ X) + Deterministic("dXX", X @ X) + text = model.str_repr(deterministic_exprs=True) + assert "dT = X.T" in text + assert "dTX = X.T @ X" in text + assert "dXX = X @ X" in text + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + assert r"{\text{X}}^{T}" in tex + assert r"{\text{X}}^{T} \cdot \text{X}" in tex + + def test_axis_permutation_shows_order(self): + with Model() as model: + w = pm.Data("w", np.zeros((2, 3, 4))) + Deterministic("d", w.dimshuffle((2, 0, 1))) + text = model.str_repr(deterministic_exprs=True) + assert "transpose(w, order=(2, 0, 1))" in text + + def test_latex_sum_axes_distinguished(self): + # Adversarial review bug 2: axis reductions must not collapse + with Model() as model: + X = pm.Data("X", np.eye(3)) + Deterministic("d", X.sum(axis=0) + X.sum(axis=1) + X.sum()) + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + assert r"\sum\_{0}\left(\text{X}\right)" in tex + assert r"\sum\_{1}\left(\text{X}\right)" in tex + assert r"\sum\left(\text{X}\right)" in tex + + def test_latex_blockwise_ops_unwrapped(self): + # Adversarial review bug 2: Blockwise core op must be shown, and + # eigh/cholesky/solve must not collapse onto the same output + with Model() as model: + X = pm.Data("X", np.eye(3)) + L = pt.linalg.cholesky(X) + sol = pt.linalg.solve(L, X) + Deterministic("d", sol[0, 0]) + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + assert r"\operatorname{Cholesky}" in tex + assert r"\operatorname{Solve}" in tex + assert "Blockwise" not in tex + + def test_latex_cast_shows_dtype(self): + with Model() as model: + s = HalfNormal("s", 1) + Deterministic("d", pt.cast(s + 1.5, "int32")) + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + assert r"\operatorname{cast}" in tex + assert r"\text{int32}" in tex + + def test_scan_renders_opaque_placeholder(self): + # Adversarial review bug 4: inner-graph machinery must not leak + from pytensor.scan import scan as pt_scan + + with Model() as model: + seq = pm.Data("seq", np.arange(5.0)) + out = pt_scan( + fn=lambda a, acc: acc + a, + sequences=seq, + outputs_info=[pt.constant(0.0, dtype=seq.dtype)], + n_steps=5, + return_updates=False, + ) + Deterministic("path", out) + text = model.str_repr(deterministic_exprs=True) + assert "AllocEmpty" not in text + assert "set_subtensor" not in text + assert "Scan{" not in text + assert "f(seq)" in text + + def test_include_params_false_takes_precedence(self): + text = self.model().str_repr(include_params=False, deterministic_exprs=True) + assert "mu = Deterministic" in text + assert "= ((alpha_a" not in text + + def test_standalone_without_named_vars_stops_at_named(self): + from pymc.printing import str_for_potential_or_deterministic + + with Model() as model: + s = HalfNormal("s", 1) + inter = s + 1 + d = Deterministic("d", s * inter) + res = str_for_potential_or_deterministic(d, deterministic_exprs=True) + assert res == "d = s * (s + 1)" + assert "~" not in res.split("=", 1)[1] + + def test_named_leaves_matched_by_identity(self): + with Model() as model: + a = Normal("a", 0, 1) + decoy = pt.vector("z") + 1 + decoy.name = "a" + Deterministic("d", a * 2) + Deterministic("e", decoy * 3) + text = model.str_repr(deterministic_exprs=True) + assert "d = a * 2" in text + # A same-named non-model variable expands instead of masquerading + assert "e = (z + 1) * 3" in text + + def test_named_leaf_beats_dict_registered_op(self): + # PPrinter checks dict-keyed registrations before condition rules; + # a potential whose owner op is a shared Elemwise instance must still + # stop traversal at the leaf boundary + from pymc.printing import str_for_potential_or_deterministic + + with Model() as model: + s = HalfNormal("s", 1) + Potential("pot", s * 2) + det = Deterministic("d", model.potentials[0] + 1) + named_vars = set(model.free_RVs) | set(model.deterministics) | set(model.potentials) + res = str_for_potential_or_deterministic( + det, named_vars=named_vars, deterministic_exprs=True + ) + assert res == "d = pot + 1" + + def test_ownerless_leaf_consistent_across_formats(self): + import pytensor + + with Model() as model: + ext = pytensor.shared(np.ones(3), name="ext") + z = Normal("z", 0, 1, shape=3) + Deterministic("d", z * ext) + text = model.str_repr(deterministic_exprs=True) + assert "d = z * ext" in text + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + assert r"\text{ext}" in tex + + +class TestDeterministicExprsParametric: + """Table-driven coverage across model families. + + Each case asserts output invariants that hold for any well-rendered body + (no leftover placeholders, no graph-internals leakage, consistent LaTeX + escaping), plus a small number of exact anchors for the ops that + distinguish the family. This catches regressions in op rendering without + snapshotting entire strings. + """ + + @staticmethod + def _cases() -> dict: + def linear_regression(m): + x = pm.Data("x", np.array([1.0, 2.0])) + a = Normal("a", 0, 1) + b = Normal("b", 0, 1) + mu = Deterministic("mu", a + b * x) + Normal("y", mu, 1) + return { + "anchors_plain": ["mu = a + (b * x)"], + "anchors_tex": [r"\text{a} + (\text{b} \cdot \text{x})"], + } + + def nonlinear(m): + s = HalfNormal("s", 1) + nl = Deterministic("nl", pt.exp(s) / (pt.sqrt(s) + pt.tanh(s))) + return { + "anchors_plain": ["Exp(s) / (Sqrt(s) + Tanh(s))"], + "anchors_tex": [r"\frac{\exp(\text{s})}{(\sqrt{\text{s}} + \tanh(\text{s}))}"], + } + + def matrix_ops(m): + X = pm.Data("X", np.eye(2)) + b = Normal("b", 0, 1, shape=2) + pm.Deterministic("p", pt.dot(X, b) + b.sum(axis=0)) + return { + "anchors_plain": ["(X @ b) + sum(b, axis=(0,))"], + "anchors_tex": [ + r"(\text{X} \cdot \text{b}) + \sum\_{0}\left(\text{b}\right)", + ], + } + + def indexing_and_slicing(m): + X = pm.Data("X", np.ones((3, 3))) + s = HalfNormal("s", 1) + pm.Deterministic("d", X[:, 0] * s + X[:2].sum(axis=0)) + return { + # plain renders real Python-style indexing... + "anchors_plain": ["X[:, 0]", "sum(X[:2]"], + # ...latex degrades gracefully to \operatorname for subtensors (#8407 open question) + "anchors_tex": [ + r"\operatorname{Subtensor}", + r"\sum\_{0}\left(\operatorname{Subtensor}", + ], + } + + def potential_only(m): + z = Normal("z") + pm.Potential("pot", -(z**2) / 2) + return { + "anchors_plain": ["pot ~ (-(z ** 2)) / 2"], + "anchors_tex": [r"\frac{(-{\text{z}}^{2})}{2}"], + } + + def hierarchical(m): + # Non-centered hierarchical model: the pooled deterministic + # stops at named parents instead of inlining their graphs + mu = Normal("mu", 0, 5) + tau = HalfNormal("tau", 5) + z = Normal("z", 0, 1, shape=8) + theta = Deterministic("theta", mu + tau * z) + Normal("y", theta, 10, shape=8) + return { + "anchors_plain": ["theta = mu + (tau * z)"], + "anchors_tex": [r"\text{mu} + (\text{tau} \cdot \text{z})"], + } + + def fixed_vs_prior_params(m): + # Parameters/hyperparameters as fixed values (scalars, and an + # inline anonymous prior) alongside named priors. Anonymous + # dists must render as distribution calls, not leak RNG internals + import pytensor + + anon = Normal.dist(0, 2) + w = pytensor.shared(np.float64(1.5)) # unnamed non-model leaf + z = Normal("z", 0, 1, shape=3) + theta = Deterministic("theta", 2.0 + 1.5 * z + w + w * 3 + anon.sum() + anon.sum()) + Normal("y", theta, 10, shape=3) + return { + "anchors_plain": [ + "2 + (1.5 * z)", + "", + "sum(Normal(0, 2), axis=None)", + ], + "anchors_tex": [ + r"(2 + (1.5 \cdot \text{z}))", + r"\text{}", + r"\sum\left(\operatorname{Normal}(0,~2)\right)", + ], + } + + return { + "linear_regression": linear_regression, + "nonlinear": nonlinear, + "matrix_ops": matrix_ops, + "indexing_and_slicing": indexing_and_slicing, + "potential_only": potential_only, + "hierarchical": hierarchical, + "fixed_vs_prior_params": fixed_vs_prior_params, + } + + @staticmethod + def _build(case_name: str) -> tuple[Model, dict]: + with Model() as model: + expected = TestDeterministicExprsParametric._cases()[case_name](model) + return model, expected + + @pytest.mark.parametrize("case_name", list(_cases()), ids=str) + def test_default_repr_uses_placeholders(self, case_name: str): + model, _ = self._build(case_name) + text = model.str_repr() + assert ("Deterministic(f(" in text) or ("Potential(f(" in text) + + @pytest.mark.parametrize("case_name", list(_cases()), ids=str) + def test_plain_bodies(self, case_name: str): + model, expected = self._build(case_name) + text = model.str_repr(deterministic_exprs=True) + # no opaque placeholders and no graph-internals leakage + assert "Deterministic(f(" not in text + assert "Potential(f(" not in text + assert "DimShuffle{" not in text + assert "ViewOp" not in text + assert "RNG(" not in text + for anchor in expected["anchors_plain"]: + assert anchor in text + + @pytest.mark.parametrize("case_name", list(_cases()), ids=str) + def test_latex_bodies(self, case_name: str): + model, expected = self._build(case_name) + tex = model.str_repr(formatting="latex", deterministic_exprs=True) + # bodies replace the wrapper operator entirely + assert r"\operatorname{Deterministic}" not in tex + # names are consistently escaped (no bare underscores anywhere) + assert "_" not in tex.replace("\\_", "") + assert "RNG(" not in tex + for anchor in expected["anchors_tex"]: + assert anchor in tex + + +class TestExpressionBounds: + """Verbosity bounds: pathological graphs degrade instead of exploding.""" + + @staticmethod + def _doubling_model(n: int) -> Model: + with Model() as model: + s = HalfNormal("s", 1) + v = s + for _ in range(n): + v = v + v + Deterministic("d", v) + return model + + def test_shared_subgraph_falls_back_to_placeholder(self): + # Each level doubles the rendered text; past the node budget the + # whole expression degrades to the opaque placeholder. + model = self._doubling_model(25) + text = model.str_repr(deterministic_exprs=True) + assert len(text) < 10_000 + assert "d = f(s)" in text + + def test_shared_subgraph_latex_bounded(self): + tex = self._doubling_model(25).str_repr(formatting="latex", deterministic_exprs=True) + assert len(tex) < 10_000 + + def test_deep_chain_degrades_gracefully(self): + # Printing must not raise RecursionError where the default repr works + with Model() as model: + s = HalfNormal("s", 1) + v = s + for _ in range(6000): + v = v + 1.0 + Deterministic("d", v) + text = model.str_repr(deterministic_exprs=True) + body = [ln for ln in text.splitlines() if ln.startswith("d =")][0] + assert "f(s)" in body # expansion stops at the depth limit + assert len(text) < 10_000 + + def test_moderate_sharing_still_expands(self): + # A realistic adstock-style chain reused twice stays fully expanded + with Model() as model: + alpha = Uniform("alpha", 0, 1) + x = Data("x", np.arange(6.0)) + acc = x + for lag in range(1, 8): + acc = acc + (alpha**lag) * x + sat = Deterministic("sat", acc / (acc + alpha)) + text = model.str_repr(deterministic_exprs=True) + start = text.index("sat = ") + end = text.find("\n", start) + end = len(text) if end == -1 else end + line = text[start:end] + assert "f(" not in line + assert line.count("alpha") >= 2 + + +class TestBodyConstants: + """Array constants describe their geometry instead of hiding behind .""" + + @staticmethod + def _model() -> Model: + with Model() as model: + x = Data("x", np.ones(4)) + vec = np.array([1.5, 2.5, 3.5]) + mat = np.arange(400, dtype="float64").reshape(20, 20) + ints = np.arange(5, dtype="int32") + a = HalfNormal("a", 1) + Deterministic("vec", pt.constant(vec) * x) + Deterministic("mat", pt.constant(mat) @ x) + Deterministic("ints", pt.constant(ints) * x) + return model + + def test_short_vector_inlines(self): + text = self._model().str_repr(deterministic_exprs=True) + assert "vec = [1.5, 2.5, 3.5] * x" in text + + def test_matrix_shows_shape(self): + text = self._model().str_repr(deterministic_exprs=True) + assert " @ x" in text + assert " * x" in text + + def test_matrix_shape_as_math_notation(self): + tex = self._model().str_repr(formatting="latex", deterministic_exprs=True) + assert r"\text{} \in \mathbb{R}^{20 \times 20}" in tex + assert r"\text{} \in \mathbb{Z}^{5}" in tex + assert r"\left[1.5,\ 2.5,\ 3.5\right]" in tex + + def test_default_repr_unchanged(self): + # Shape-aware rendering is opt-in only; defaults keep + text = self._model().str_repr() + assert "(20, 20)" not in text