docs: guide to PyTensor compilation for PyMC users [first draft / RFC]

[drbenvincent]

⚠️ Status: first draft — seeking high-level feedback

Putting this up early to get first impressions and direction before I invest in polishing. Feedback on scope, framing, and accuracy is more useful than wording at this stage; I'm happy to restructure significantly.

Summary

A new user-guide page, "Understanding PyTensor compilation (a guide for PyMC users)", for people who reach PyTensor through PyMC rather than using it directly. It builds a mental model of the compilation pipeline and answers two recurring PyMC questions: "why does sampling take a while to start?" (compilation) and "why is the first run slow but later runs faster?" (caching).

What's included

• A Mermaid system diagram of the full pipeline (PyMC layer → graph prep → Stage 1 rewriting → backend selection → Stage 2 linking → sampler runtime), with the PyTensor/​sampler boundary made explicit.
• Stage 1 (rewriting) explained pass by pass — merge/useless, canonicalize, stabilize, specialize/fusion, backend-specific + in-place — in plain language, with collapsible "Under the hood" notes linking to the implementing modules/functions.
• Stage 2 (linking) explained backend by backend — Numba (default), C/CVM, JAX, PyTorch, MLX, Python VM.
• How the backend is chosen (backend= / compile_kwargs / config.linker) vs. how the sampler is chosen (nuts_sampler=: pymc/nutpie/numpyro/blackjax), and why they're mostly-independent choices.
• Toctree entry plus the doc-build deps (sphinx-design, sphinxcontrib-mermaid) and conf.py mermaid config.

Open questions for reviewers

• Is the user guide the right home, and is the "opinionated, PyMC-first" framing acceptable for the docs?
• Is the level of internal detail right (the folds reference specific modules/functions/positions)? Too much, too little?
• Anything technically off — especially the Numba-default story, the sampler/backend coupling, and the caching / "Stage 1 isn't cached" note?

Notes

• Docs-only change; no code paths touched.
• Builds locally with python -m sphinx -b html ./doc ./html; cross-references (FunctionGraph, Mode, jax.jit, torch.compile) resolve and the Mermaid diagram + collapsible notes render.

Made with Cursor

[ricardoV94]

sounds like a PyMC documentation not PyTensor. There's already one doc answering the "graph" approach part https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/pymc_pytensor.html

[drbenvincent]

Thanks @ricardoV94 — agreed, this belongs in the PyMC docs alongside the existing PyMC and PyTensor notebook rather than here.

I've moved it over and reframed it as a companion to that notebook (placed right after it in the core-notebooks toctree): pymc-devs/pymc#8326. Closing this one in favour of that PR.