ENH: HTML model table with LaTeX-rendered equations for notebook display

[drbenvincent]

Context

PR #382 added model_table() to pymc_extras.printing, which renders a PyMC model as a text table via the Rich library. The output looks like:

 Variable  Expression         Dimensions
─────────────────────────────────────────────────────
      x =  Data               subject[20] × param[2]
      y =  Data               subject[20]

   beta ~  Normal(0, 1)       param[2]
  sigma ~  HalfNormal(0, 1)
                              Parameter count = 3

     mu =  f(beta)            subject[20]

  y_obs ~  Normal(mu, sigma)  subject[20]

This is useful, but in notebook environments (Jupyter, VS Code, Quarto) users expect rich HTML output with properly rendered mathematics. PyMC already has the building blocks: str_for_dist and str_for_potential_or_deterministic in pymc.printing both accept a formatting="latex" argument and produce valid LaTeX strings.

During PR #382 review, @jessegrabowski suggested splitting the string-building logic from the presentation logic so the intermediate representation could feed into different renderers — and @ricardoV94 agreed and refactored accordingly. The utilities variable_expression and dims_expression in pymc_extras.printing now return plain strings independent of any output format, which sets the stage for an HTML backend.

Goal

Add an HTML rendering path for model_table that:

1. Returns an HTML string (or an object with _repr_html_) suitable for automatic rich display in Jupyter / IPython.
2. Renders distribution expressions as LaTeX math (using MathJax / KaTeX, which Jupyter already provides).
3. Preserves the same grouped structure (data, free RVs, deterministics, potentials, observed RVs) and parameter count as the existing Rich table.

Possible approach

1. Produce LaTeX expressions from the existing utilities

pymc.printing.str_for_dist(var, formatting="latex") already yields strings like $\text{beta} \sim \text{Normal}(0,~1)$. The variable_expression helper in pymc_extras.printing currently calls str_for_dist with the default formatting="plain". Passing formatting="latex" instead (or adding a formatting parameter to variable_expression) would give us LaTeX-ready expression strings.

2. Build an HTML table

Construct a simple <table> (or use a lightweight helper) with three columns — Variable, Expression, Dimensions — mirroring the current Rich table. Wrap each expression cell's content in $...$ so that MathJax/KaTeX renders it automatically when displayed in a notebook.

A minimal sketch:

def model_table_html(model, **kwargs) -> str:
    """Return an HTML string summarising the model, with LaTeX-rendered expressions."""
    rows = _collect_rows(model, formatting="latex", **kwargs)  # reuse extraction logic
    html = ['<table class="pymc-model-table">']
    html.append("<thead><tr><th>Variable</th><th>Expression</th><th>Dimensions</th></tr></thead>")
    html.append("<tbody>")
    for row in rows:
        html.append(f"<tr><td>{row.name}</td><td>{row.expression}</td><td>{row.dims}</td></tr>")
    html.append("</tbody></table>")
    return "\n".join(html)

3. Wrap in a displayable object

Return a small wrapper class with a _repr_html_ method so that Jupyter displays it automatically:

class ModelTable:
    def __init__(self, html: str, rich_table: Table):
        self._html = html
        self._rich = rich_table

    def _repr_html_(self):
        return self._html

    def __rich_console__(self, console, options):
        yield self._rich

This way the same object works in plain terminals (via Rich) and in notebooks (via _repr_html_).

4. Consider styling

A small embedded <style> block could improve readability: right-align the variable column, add light group separators between variable groups, and use a monospace font for dimension annotations.

Open questions

• Where should this live? Adding it to pymc_extras.printing alongside model_table seems natural. It could be a format parameter on model_table itself (e.g. format="html") or a separate function.
• Should the Model get a _repr_html_? PyMC's Model.__init__ already wires up _repr_latex_ via str_for_model. A _repr_html_ that calls the new HTML table would give notebooks a richer default display, but that change would live in PyMC core, not pymc-extras.
• Deterministic expressions in LaTeX. The current variable_expression shows deterministics as f(a, b, ...). In LaTeX mode, it might be worth rendering the actual symbolic expression where feasible (e.g. x \cdot \beta instead of f(\text{beta})), but this is a deeper effort and could be a follow-up.

[ricardoV94]

rich can already export to html, it's in the docstrings example. Not sure how good/bad that is:

from rich.console import Console

console = Console(record=True)
console.print(table)   # <- output of the model_table()
text_export = console.export_text()
html_export = console.export_html()
svg_export = console.export_svg()