From ed15d8b9c8d389c5a987ee3825a6e1a1ad98e345 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Sat, 13 Jun 2026 21:29:07 +0100 Subject: [PATCH 1/4] docs: add PyTensor compilation guide for PyMC users Add an opinionated User Guide page explaining PyTensor's two-stage compilation pipeline (graph rewriting then backend linking) from the perspective of a PyMC user, including a Mermaid system diagram covering the Numba/C/JAX/PyTorch/MLX/VM backends and the runtime logp+gradient callable consumed by samplers (PyMC NUTS, nutpie). Enable the sphinxcontrib-mermaid extension (and render plain ```mermaid fences via myst_fence_as_directive) so the diagram renders in the docs build, on GitHub, and in editor previews. Co-authored-by: Cursor --- doc/compilation_for_pymc_users.md | 171 ++++++++++++++++++++++++++++++ doc/conf.py | 5 + doc/environment.yml | 1 + doc/user_guide.rst | 1 + 4 files changed, 178 insertions(+) create mode 100644 doc/compilation_for_pymc_users.md diff --git a/doc/compilation_for_pymc_users.md b/doc/compilation_for_pymc_users.md new file mode 100644 index 0000000000..6a6e625674 --- /dev/null +++ b/doc/compilation_for_pymc_users.md @@ -0,0 +1,171 @@ +(compilation_for_pymc_users)= +# Understanding PyTensor compilation (a guide for PyMC users) + +:::{note} +This page is deliberately *opinionated* and slightly broader in scope than the rest of the documentation: it explains PyTensor from the vantage point of someone who arrives at it **through PyMC** rather than using PyTensor directly. That framing is intentional — a large share of PyTensor users meet it at the top of the stack, via PyMC — but it does mean this guide makes some simplifications and value judgements that a pure-PyTensor reference would not. +::: + +If you build models with PyMC, you don't usually call PyTensor yourself. But PyTensor is doing a lot of work on your behalf every time you sample, and understanding *what* it does helps explain two things PyMC users often wonder about: + +- **Why does sampling sometimes take a while to even start?** (That's compilation.) +- **Why is the first run slow but later runs faster?** (That's caching.) + +This guide gives you a mental model of the pipeline and the choices available, without assuming you've ever written a line of PyTensor. + +## The one-sentence mental model + +When PyMC samples your model, it asks PyTensor to turn the model's math (the log-probability and its gradient) into a fast, callable function. PyTensor does this in **two stages**: it first **rewrites** the symbolic math into a better form, then **links** that form into something executable using a chosen **backend**. + +## The pipeline at a glance + + + +```mermaid +flowchart TD + %% ============ PyMC layer ============ + subgraph PYMC["PyMC layer"] + A["PyMC model
(pm.Model: priors, likelihood)"] + A --> B["Build symbolic graph
logp / dlogp / random fns"] + B --> C["pytensor.function(inputs, outputs, mode=...)"] + end + + %% ============ Graph prep ============ + subgraph PREP["Graph preparation (pure Python)"] + C --> D["clone graph, apply givens/updates,
collect shared variables"] + D --> E["wrap in FunctionGraph"] + end + + %% ============ Stage 1: rewriting ============ + subgraph REWRITE["STAGE 1 — Graph rewriting (pure Python, backend-agnostic)"] + E --> F["run optimizer pipeline (optdb)"] + F --> G["merge / useless"] + G --> H["canonicalize
(equilibrium, fixed-point loop)"] + H --> I["stabilize
numerically-stable forms"] + I --> J["specialize + fusion"] + J --> K["backend-specific rewrites
+ inplace / destroy handler"] + K --> L["Rewritten FunctionGraph"] + end + + %% ============ Backend selection ============ + L --> M{"Which linker?
(from mode / config.linker)"} + M -->|"'auto' = DEFAULT"| NUMBA + M -->|"c / cvm (opt-in)"| CBACK + M -->|"jax"| JAXB + M -->|"pytorch"| PTB + M -->|"mlx (Apple GPU)"| MLXB + M -->|"py / vm"| PYB + + %% ============ Stage 2: linking ============ + subgraph LINK["STAGE 2 — Linking (backend-specific)"] + subgraph NUMBA["Numba backend (default, pip-friendly)"] + N1["graph -> Python"] + N1 --> N2["numba.njit JIT via LLVM"] + N2 --> N3["on-disk cache
(numba__cache=True)"] + end + subgraph CBACK["C / CVM backend (needs C++ compiler)"] + C1["generate C++ source"] + C1 --> C2["g++/clang++ subprocess"] + C2 --> C3["cached .so modules"] + end + subgraph JAXB["JAX backend"] + J1["graph -> JAX"] + J1 --> J2["jax.jit -> XLA"] + end + subgraph PTB["PyTorch backend"] + T1["graph -> PyTorch"] + T1 --> T2["torch.compile
(TorchDynamo/Inductor)"] + end + subgraph MLXB["MLX backend (Apple Silicon GPU)"] + X1["graph -> MLX"] + X1 --> X2["mx.compile (lazy JIT)"] + end + subgraph PYB["Python VM backend"] + P1["Op.perform in pure Python"] + end + end + + %% ============ Compiled artifact (bridges linking and runtime) ============ + N3 --> Z["Compiled callable
params θ (point in parameter space) → logp(θ), ∇logp(θ)
(observed data baked in as constants)"] + C3 --> Z + J2 --> Z + T2 --> Z + X2 --> Z + P1 --> Z + + %% ============ Runtime ============ + Z --> R + subgraph RUN["Runtime — MCMC sampling (function called many times)"] + R["Sampler loop:
propose new θ → evaluate logp + gradient → repeat"] + R --> S["Sampler engine, e.g.
PyMC NUTS (Python)
or nutpie / nuts-rs (Rust)"] + end +``` + +## Walking through the stages + +### Graph preparation + +PyMC hands PyTensor a symbolic graph — a recipe of mathematical operations, not yet any numbers. PyTensor clones it, resolves shared variables and updates, and wraps it in a `FunctionGraph`. This is bookkeeping; it's cheap. + +### Stage 1 — rewriting (the "optimizer") + +PyTensor rewrites the graph to be faster and more numerically stable. This is where, for example, a naive `log(1 + exp(x))` becomes a stable `softplus`, duplicate sub-expressions get merged, and element-wise operations get fused together. The heavy passes (`canonicalize`, `stabilize`, `specialize`) run to a *fixed point* — they keep applying rewrites until the graph stops changing. + +Two things are worth internalising: + +- **This stage is pure Python**, and its cost scales with the size of your graph. Large hierarchical models have large graphs, so this can be a real chunk of "time before sampling". +- **This stage is backend-agnostic.** Whether you end up on Numba, C, JAX, PyTorch, or MLX, the same rewriting happens first. + +### Stage 2 — linking (choosing a backend) + +The rewritten graph is turned into an executable. *How* depends on the linker/backend: + +| Backend | How it compiles | Needs a system compiler? | Notes | +|---|---|---|---| +| **Numba** *(default)* | Converts the graph to Python, then JIT-compiles via LLVM | No (ships as wheels) | This is why `pip install pymc` "just works" without conda. First compile is slow; results are cached on disk. | +| **C / CVM** | Generates C++ source and compiles `.so` files with `g++`/`clang++` | Yes | Historically the default. Strong on-disk module cache. Now opt-in. | +| **JAX** | Converts to JAX, compiles via `jax.jit` → XLA | No | Great for GPU/TPU and `vmap`-style workflows. | +| **PyTorch** | Converts to PyTorch, compiles via `torch.compile` (TorchDynamo/Inductor) | No | Runs on PyTorch's CPU/GPU devices. | +| **MLX** | Converts to MLX, compiles via `mx.compile` | No | Targets Apple Silicon GPUs. | +| **Python VM** | Runs each operation's pure-Python implementation | No | Slowest at runtime, fastest to "compile". Handy for debugging. | + +The default backend resolves from `config.linker = "auto"`, which currently maps to **Numba**. + +### Runtime + +Functionally, the compiled callable is a map from a **point in parameter space** to the model's **log-probability and its gradient**: `θ → (logp(θ), ∇logp(θ))`. The observed data isn't an argument — it's baked into the function as constants when the function is built. (PyMC actually compiles a few such functions: the log-density and its gradient for sampling, plus functions for drawing from priors/posterior predictive.) + +An MCMC sampler calls this function thousands of times, proposing new parameter points and reading back `logp` and its gradient to decide where to go next. Time spent *here* is sampling time, which is separate from the compilation time described above — speeding up compilation does not speed up sampling, and vice versa. + +The sampler engine that drives those calls is a separate choice from the PyTensor backend. It can be PyMC's built-in NUTS (Python), or an external sampler such as [nutpie](https://github.com/pymc-devs/nutpie), which runs the NUTS loop in Rust (via `nuts-rs`) while still calling a PyTensor-compiled logp/gradient under the hood. The Rust there is the *sampler*, not a PyTensor compilation backend — a useful distinction, since it's easy to assume "the fast Rust thing" is doing the compilation when in fact PyTensor still builds the function it evaluates. + +## Caching: why the second run is faster + +PyTensor caches the **linked artifact** so repeated builds can skip recompilation: + +- **Numba** keeps an on-disk cache (enabled by default via `numba__cache`). +- **C / CVM** caches compiled `.so` modules in your compile directory. + +Note a current limitation that matters in practice: the cache is for the *linked* output. Stage 1 rewriting is **not** cached and is re-run on every build, even for a structurally identical model. + +## Practical tips for PyMC users + +- **Iterating on model structure? Compile faster, run slower.** Use the `FAST_COMPILE` mode, which skips many rewrites and uses the pure-Python VM (no C/LLVM compilation at all). In PyMC you can often pass this via `compile_kwargs={"mode": "FAST_COMPILE"}`. Switch back to the default for the real, long sampling run. +- **Want to know where the time goes?** Turn on profiling: + + ```python + import pytensor + pytensor.config.profile = True # per-function rewrite vs link time + pytensor.config.profile_optimizer = True # per-rewrite breakdown + ``` + +- **"It recompiles every morning."** That usually means the on-disk cache isn't being hit across sessions. Check that your compile directory is stable and not being cleared between runs. +- **Choosing a backend.** Stick with the default (Numba) unless you have a specific reason: JAX for GPU/TPU, PyTorch to integrate with the PyTorch ecosystem, MLX for Apple Silicon GPUs, C/CVM if you specifically need the C runtime characteristics. + +## Where to go next + +- {ref}`optimizations` — the catalogue of graph rewrites applied in Stage 1. +- {ref}`using_modes` — modes and linkers in more depth. +- {doc}`library/compile/mode` — the `Mode` / linker API reference. diff --git a/doc/conf.py b/doc/conf.py index e10dcffb90..22147464d7 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -25,6 +25,7 @@ "myst_nb", "generate_gallery", "sphinx_sitemap", + "sphinxcontrib.mermaid", ] # Don't auto-generate summary for class members. @@ -329,6 +330,10 @@ def fallback_source(): ] myst_dmath_double_inline = True +# Render plain ```mermaid fences as the mermaid directive, so the same source +# renders in the Sphinx build, on GitHub, and in editor markdown previews. +myst_fence_as_directive = ["mermaid"] + citation_code = f""" ```bibtex @incollection{{citekey, diff --git a/doc/environment.yml b/doc/environment.yml index 8d3419e7a7..17dc05c423 100644 --- a/doc/environment.yml +++ b/doc/environment.yml @@ -16,6 +16,7 @@ dependencies: - sphinx-copybutton - sphinx-design - sphinx-sitemap + - sphinxcontrib-mermaid - pygments - pydot - ipython diff --git a/doc/user_guide.rst b/doc/user_guide.rst index 26401da596..cf80b72836 100644 --- a/doc/user_guide.rst +++ b/doc/user_guide.rst @@ -8,6 +8,7 @@ User Guide .. toctree:: :maxdepth: 1 + compilation_for_pymc_users extending/index optimizations troubleshooting From a8362db212a59e1683c89d2503dbfe8f2b2de20e Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Sat, 13 Jun 2026 21:53:20 +0100 Subject: [PATCH 2/4] better --- doc/compilation_for_pymc_users.md | 107 +++++++++++++++++++++++++----- 1 file changed, 89 insertions(+), 18 deletions(-) diff --git a/doc/compilation_for_pymc_users.md b/doc/compilation_for_pymc_users.md index 6a6e625674..3a885043ec 100644 --- a/doc/compilation_for_pymc_users.md +++ b/doc/compilation_for_pymc_users.md @@ -19,8 +19,20 @@ When PyMC samples your model, it asks PyTensor to turn the model's math (the log ## The pipeline at a glance ```mermaid @@ -50,13 +62,13 @@ flowchart TD end %% ============ Backend selection ============ - L --> M{"Which linker?
(from mode / config.linker)"} - M -->|"'auto' = DEFAULT"| NUMBA - M -->|"c / cvm (opt-in)"| CBACK - M -->|"jax"| JAXB - M -->|"pytorch"| PTB - M -->|"mlx (Apple GPU)"| MLXB - M -->|"py / vm"| PYB + L --> M{"Which backend (linker)?
pm.sample(backend=...)
or compile_kwargs mode=... / config.linker"} + M -->|"'numba' (default)"| NUMBA + M -->|"'c' (→ cvm)"| CBACK + M -->|"'jax'"| JAXB + M -->|"'pytorch'"| PTB + M -->|"'mlx' (Apple GPU)"| MLXB + M -->|"'py' / 'vm'"| PYB %% ============ Stage 2: linking ============ subgraph LINK["STAGE 2 — Linking (backend-specific)"] @@ -87,20 +99,34 @@ flowchart TD end end - %% ============ Compiled artifact (bridges linking and runtime) ============ - N3 --> Z["Compiled callable
params θ (point in parameter space) → logp(θ), ∇logp(θ)
(observed data baked in as constants)"] + %% ============ Compiled artifact = PyTensor's final output (the boundary) ============ + N3 --> Z["Compiled callable — PyTensor's output
θ (a point in parameter space) → logp(θ), ∇logp(θ)
(observed data baked in as constants)"] C3 --> Z J2 --> Z T2 --> Z X2 --> Z P1 --> Z - %% ============ Runtime ============ - Z --> R - subgraph RUN["Runtime — MCMC sampling (function called many times)"] - R["Sampler loop:
propose new θ → evaluate logp + gradient → repeat"] - R --> S["Sampler engine, e.g.
PyMC NUTS (Python)
or nutpie / nuts-rs (Rust)"] + %% ============ Runtime: the sampler engine (NOT PyTensor) ============ + subgraph RUN["Runtime — MCMC sampling (the sampler engine, NOT PyTensor)"] + direction TB + LOOP["Sampler loop:
propose θ → call the compiled logp/∇logp → accept/reject → repeat (many times)"] + LOOP --> ENG{"Which sampler engine?
pm.sample(nuts_sampler=...)"} + ENG -->|"'pymc'"| SP1["PyMC NUTS
pure-Python loop
(calls a callable from any CPU backend, e.g. C / Numba)"] + ENG -->|"'nutpie' (default if installed)"| SP2["nutpie / nuts-rs
Rust loop
(needs a Numba or JAX callable)"] + ENG -->|"'numpyro'"| SP3["NumPyro NUTS
JAX loop
(needs a JAX callable)"] + ENG -->|"'blackjax'"| SP4["BlackJAX NUTS
JAX loop
(needs a JAX callable)"] end + + %% Connect the boundary artifact into the runtime AFTER the subgraph is declared, + %% targeting the already-declared LOOP node so Mermaid links into the cluster. + Z --> LOOP + + classDef sampler fill:#fff3e0,stroke:#e65100,color:#000; + class LOOP,SP1,SP2,SP3,SP4 sampler; + + classDef choice fill:#e3f2fd,stroke:#1565c0,color:#000; + class M,ENG choice; ``` ## Walking through the stages @@ -133,13 +159,58 @@ The rewritten graph is turned into an executable. *How* depends on the linker/ba The default backend resolves from `config.linker = "auto"`, which currently maps to **Numba**. +#### How the backend gets chosen + +Just as `nuts_sampler=` picks the sampler engine, a different knob picks the PyTensor backend. From PyMC, the simplest way is the `backend=` argument to `pm.sample`: + +```python +import pymc as pm + +with model: + # Default-ish: Numba. Other common choices: "c", "jax". + idata = pm.sample(backend="numba") +``` + +Under the hood, PyMC turns `backend=` into a PyTensor *compilation mode* (note that `backend="c"` maps to PyTensor's combined C+VM mode, `"cvm"`). If you need finer control you can pass a mode directly via `compile_kwargs` instead — but you can only set one of the two: + +```python +with model: + # Equivalent to backend="jax"; do NOT also pass backend=... + idata = pm.sample(compile_kwargs={"mode": "JAX"}) +``` + +If you're calling PyTensor directly (no PyMC), the same choice is made with `config.linker` for the whole session, or per call via the `mode` argument to `pytensor.function`: + +```python +import pytensor + +# Session-wide default for every compiled function +pytensor.config.linker = "jax" + +# Or per function, overriding the session default +f = pytensor.function([x], y, mode="NUMBA") +``` + +The two choices — backend and sampler — are made at different points in the pipeline (see the two blue decision diamonds in the diagram), and as noted below they're mostly independent. The exception is that JAX-based samplers force the JAX backend. + ### Runtime Functionally, the compiled callable is a map from a **point in parameter space** to the model's **log-probability and its gradient**: `θ → (logp(θ), ∇logp(θ))`. The observed data isn't an argument — it's baked into the function as constants when the function is built. (PyMC actually compiles a few such functions: the log-density and its gradient for sampling, plus functions for drawing from priors/posterior predictive.) -An MCMC sampler calls this function thousands of times, proposing new parameter points and reading back `logp` and its gradient to decide where to go next. Time spent *here* is sampling time, which is separate from the compilation time described above — speeding up compilation does not speed up sampling, and vice versa. +An MCMC sampler calls this function many times, proposing new parameter points and reading back `logp` and its gradient to decide where to go next. Time spent *here* is sampling time, which is separate from the compilation time described above — speeding up compilation does not speed up sampling, and vice versa. + +**This is the key boundary.** Everything above the compiled callable in the diagram is *PyTensor* turning your model's math into a function. Everything below it is a *sampler engine* repeatedly calling that function. The sampler is a separate piece of software with its own loop; PyTensor's job is finished once the callable exists. + +You pick the engine with `pm.sample(nuts_sampler=...)`. The options are: + +- **`"pymc"`** — PyMC's built-in NUTS. The loop is pure Python and calls a PyTensor callable from any CPU backend (typically C or Numba). +- **`"nutpie"`** — runs the NUTS loop in Rust (via [`nuts-rs`](https://github.com/pymc-devs/nutpie)). It calls a PyTensor callable compiled to either the **Numba** or the **JAX** backend. This is the default when nutpie is installed. +- **`"numpyro"`** — NumPyro's NUTS, whose loop runs in **JAX**. It requires the **JAX** backend. +- **`"blackjax"`** — BlackJAX's NUTS, also a **JAX** loop requiring the **JAX** backend. + +So the engine and the PyTensor backend are *mostly* an independent choice, with one important coupling: the JAX-based engines (`numpyro`, `blackjax`, and `nutpie` in JAX mode) need PyTensor to have compiled the callable to JAX. PyMC handles wiring this up for you. -The sampler engine that drives those calls is a separate choice from the PyTensor backend. It can be PyMC's built-in NUTS (Python), or an external sampler such as [nutpie](https://github.com/pymc-devs/nutpie), which runs the NUTS loop in Rust (via `nuts-rs`) while still calling a PyTensor-compiled logp/gradient under the hood. The Rust there is the *sampler*, not a PyTensor compilation backend — a useful distinction, since it's easy to assume "the fast Rust thing" is doing the compilation when in fact PyTensor still builds the function it evaluates. +A common misconception is worth heading off: with nutpie it's easy to assume "the fast Rust thing" is doing the compilation. It isn't — the Rust is the *sampler loop*, not a PyTensor compilation backend. PyTensor still builds the logp/gradient function that the Rust loop evaluates. ## Caching: why the second run is faster From 9f1071e5c812401327f8b9994ff6c2534ee7400b Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Sat, 13 Jun 2026 22:04:55 +0100 Subject: [PATCH 3/4] more detail on pytensor --- doc/compilation_for_pymc_users.md | 63 +++++++++++++++++++++++++++++-- 1 file changed, 59 insertions(+), 4 deletions(-) diff --git a/doc/compilation_for_pymc_users.md b/doc/compilation_for_pymc_users.md index 3a885043ec..3626f5724e 100644 --- a/doc/compilation_for_pymc_users.md +++ b/doc/compilation_for_pymc_users.md @@ -137,12 +137,39 @@ PyMC hands PyTensor a symbolic graph — a recipe of mathematical operations, no ### Stage 1 — rewriting (the "optimizer") -PyTensor rewrites the graph to be faster and more numerically stable. This is where, for example, a naive `log(1 + exp(x))` becomes a stable `softplus`, duplicate sub-expressions get merged, and element-wise operations get fused together. The heavy passes (`canonicalize`, `stabilize`, `specialize`) run to a *fixed point* — they keep applying rewrites until the graph stops changing. +PyTensor rewrites the graph to be faster and more numerically stable. This is where, for example, a naive `log(1 + x)` becomes a stable `log1p(x)`, duplicate sub-expressions get merged, and element-wise operations get fused together. Two things are worth internalising up front: **this stage is pure Python**, and its cost scales with the size of your graph (large hierarchical models have large graphs, so this can be a real chunk of "time before sampling"); and **this stage is backend-agnostic** — whether you end up on Numba, C, JAX, PyTorch, or MLX, almost all of the same rewriting happens first. -Two things are worth internalising: +The subsections below walk through the boxes in the diagram's "Stage 1" lane, in order. Each names the module or function that implements it, so you can jump into the source. -- **This stage is pure Python**, and its cost scales with the size of your graph. Large hierarchical models have large graphs, so this can be a real chunk of "time before sampling". -- **This stage is backend-agnostic.** Whether you end up on Numba, C, JAX, PyTorch, or MLX, the same rewriting happens first. +#### The optimizer pipeline (`optdb`) + +Every rewrite lives in one big, ordered registry: `pytensor.compile.mode.optdb`, an instance of `pytensor.graph.rewriting.db.SequenceDB`. Entries are registered at a numeric `position` (merge at `0`, `canonicalize` at `1`, `stabilize` at `1.5`, `specialize` at `2`, fusion at `49`, inplace at `49.5`+), and the pipeline runs them in that order. Each entry also carries *tags* such as `"fast_run"`, `"fast_compile"`, `"numba"`, or `"inplace"`. + +Which entries actually run is decided by the {class}`pytensor.compile.mode.Mode` you compile with. A `Mode` pairs a *linker* (Stage 2) with an *optimizer query* (`pytensor.graph.rewriting.db.RewriteDatabaseQuery`) that selects rewrites by tag. For instance the default Numba mode queries `include=["fast_run", "numba"]`, while `FAST_COMPILE` queries `include=["fast_compile", "py_only"]` and so skips most of the pipeline. This is exactly the knob `FAST_COMPILE` and friends are turning. + +#### `merge` / `useless` + +The first passes are cheap clean-ups. **Merge** (`pytensor.graph.rewriting.basic.MergeOptimizer`, registered as `merge1`, `merge1.1`, `merge2`, `merge3` at several positions) collapses identical sub-graphs so a repeated expression is computed once. **Useless** (the `"useless"` entry at position `0.6`, a `pytensor.graph.rewriting.db.TopoDB` over a `LocalGroupDB` of small `local_useless_*` node rewriters) strips nodes that don't affect the output, e.g. an identity reshape or a fill that broadcasts to a shape already in hand. + +#### `canonicalize` (fixed-point loop) + +**Canonicalize** (position `1`) rewrites the graph into a single normal form so that later passes only have to recognise one shape of each pattern. It is an `pytensor.graph.rewriting.db.EquilibriumDB`: it applies all of its member rewrites repeatedly until the graph stops changing (reaches *equilibrium* / a fixed point). Developers add a local rewrite here with the `@register_canonicalize` decorator from `pytensor.tensor.rewriting.basic`. Typical work: reordering and flattening `add`/`mul` chains, folding constants, and pushing operations into a canonical position. + +#### `stabilize` (numerically-stable forms) + +**Stabilize** (position `1.5`, also an `EquilibriumDB`, populated via `@register_stabilize`) swaps numerically fragile expressions for equivalent stable ones — the step PyMC users feel most, because log-probabilities are full of logs and exps. Concrete examples from `pytensor.tensor.rewriting.math`: `local_log1p` turns `log(1 + x)` into `log1p(x)`, `local_expm1` turns `exp(x) - 1` into `expm1(x)`, and `local_log1p_plusminus_exp` turns `log1p(exp(x))` into `log1pexp(x)` (the softplus). These avoid catastrophic cancellation and overflow that the naive forms suffer near the extremes. + +#### `specialize` + `fusion` + +**Specialize** (position `2`, `EquilibriumDB`, `@register_specialize`) replaces general operations with faster special-cased ones (for example rewriting `x ** 2` or `sum(x**2)` into cheaper equivalents). **Fusion** then merges many element-wise operations into a single `Composite` op so the runtime walks the data once instead of allocating a temporary per operation: see the `"elemwise_fusion"` sequence (the `FusionOptimizer`) and `"add_mul_fusion"` in `pytensor.tensor.rewriting.elemwise`. Fusion is tagged `"fusion"`, which is why `FAST_COMPILE` (and JAX, which does its own fusion) leave it out. + +#### Backend-specific rewrites + inplace / destroy handler + +Two kinds of rewrites run near the end. **Backend-specific** rewrites are pulled in by the mode's tag query — the Numba mode adds `"numba"`-tagged rewrites (`pytensor.tensor.rewriting.numba`), the JAX mode adds `"jax"`-tagged ones (`pytensor.tensor.rewriting.jax`), and so on; conversely C-only ops and BLAS rewrites are *excluded* for the array-library backends. **Inplace** rewrites (tagged `"inplace"`) let ops write into their inputs' memory to avoid allocations: `AddDestroyHandler` (position `49.5`) first attaches a `pytensor.graph.destroyhandler.DestroyHandler` that tracks which buffers may be clobbered, then rewrites like `InplaceElemwiseOptimizer` (`pytensor.tensor.rewriting.elemwise`, position `50.5`) switch ops to their in-place variants. JAX excludes inplace entirely, since it manages memory itself. + +#### Rewritten FunctionGraph + +The output of Stage 1 is a new {class}`pytensor.graph.fg.FunctionGraph` — the same inputs and outputs, but a leaner, more stable, backend-tailored computation. Nothing has been compiled yet; that is Stage 2's job. (To inspect this graph without compiling a full function, apply the rewrites directly with `pytensor.graph.rewriting.utils.rewrite_graph(y, include=("canonicalize", "specialize"))` and print the result with `pytensor.dprint`.) ### Stage 2 — linking (choosing a backend) @@ -193,6 +220,34 @@ f = pytensor.function([x], y, mode="NUMBA") The two choices — backend and sampler — are made at different points in the pipeline (see the two blue decision diamonds in the diagram), and as noted below they're mostly independent. The exception is that JAX-based samplers force the JAX backend. +#### How linking works (the shared mechanism) + +Every backend is implemented as a *linker*, a subclass of `pytensor.link.basic.Linker`; the {class}`pytensor.compile.mode.Mode`'s linker is what Stage 2 runs. Most linkers share one trick: a single-dispatch registry that translates each PyTensor `Op` into the target language. Numba uses `pytensor.link.numba.dispatch.numba_funcify`, JAX uses `pytensor.link.jax.dispatch.jax_funcify`, PyTorch uses `pytensor.link.pytorch.dispatch.pytorch_funcify`, and MLX uses `pytensor.link.mlx.dispatch.mlx_funcify`. To add backend support for a new `Op`, you register a function with the matching dispatcher. The C and Python backends are the exception: rather than a `*_funcify` registry, they ask each `Op` directly for code through `Op.c_code` or run its `Op.perform` method. The subsections below cover each branch of the diagram's "Stage 2" lane. + +#### Numba (default) + +`pytensor.link.numba.linker.NumbaLinker`. `numba_funcify` walks the rewritten `FunctionGraph` and emits one Python function that wires together each op's Numba implementation; `pytensor.link.numba.dispatch.numba_njit` then JIT-compiles it with [`numba.njit`](https://numba.readthedocs.io/en/stable/user/jit.html) via LLVM. The first call pays the compilation cost; afterwards it runs as native machine code. With `config.numba__cache = True` (the default) the compiled result is written to an on-disk cache (keyed per op, see `numba_funcify_and_cache_key`), so a later session can skip recompiling. Numba ships as wheels, which is why `pip install pymc` "just works" without a system compiler. + +#### C / CVM + +`pytensor.link.c.basic.CLinker` generates C++ for each op from its `Op.c_code`, compiles a `.so` with the system `g++`/`clang++`, and caches the module in your compile directory. The default `"cvm"` linker is `pytensor.link.vm.VMLinker` with `use_cloop=True`: a hybrid where a fast C "virtual machine" loop steps through op *thunks*, each of which may be C-compiled or fall back to Python. Pure `"c"` (`CLinker` / `OpWiseCLinker`) compiles as much of the graph into C as possible. This family was historically the default and has a strong on-disk module cache, but it needs a working C++ toolchain — hence the move to Numba for the pip story. + +#### JAX + +`pytensor.link.jax.linker.JAXLinker`. `jax_funcify` translates the graph into a JAX-traceable Python function, which the linker hands to {func}`jax.jit` → XLA. This is the GPU/TPU-friendly path, and it's also the backend the JAX samplers (`numpyro`, `blackjax`, and `nutpie` in JAX mode) consume directly. + +#### PyTorch + +`pytensor.link.pytorch.linker.PytorchLinker`. `pytorch_funcify` builds a PyTorch function, compiled with {func}`torch.compile` (TorchDynamo + Inductor), running on PyTorch's CPU/GPU devices. + +#### MLX + +`pytensor.link.mlx.linker.MLXLinker`. `mlx_funcify` builds an MLX function compiled lazily with `mx.compile`, targeting Apple Silicon GPUs. + +#### Python VM + +The reference backend: `pytensor.link.basic.PerformLinker` (or `pytensor.link.vm.VMLinker` with `use_cloop=False`) steps through the graph calling each op's `Op.perform` in pure Python. There is essentially no compilation step, so it is the fastest to "build" and the slowest to run — ideal for debugging, and what `FAST_COMPILE` uses. + ### Runtime Functionally, the compiled callable is a map from a **point in parameter space** to the model's **log-probability and its gradient**: `θ → (logp(θ), ∇logp(θ))`. The observed data isn't an argument — it's baked into the function as constants when the function is built. (PyMC actually compiles a few such functions: the log-density and its gradient for sampling, plus functions for drawing from priors/posterior predictive.) From 9cdc6e0e4e68ce9ccfb9270a006da5d1f3135a31 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Sat, 13 Jun 2026 22:30:21 +0100 Subject: [PATCH 4/4] docs: rewrite compilation guide for approachability Rework the PyMC-user compilation guide to be readable without prior PyTensor knowledge: explain each Stage 1 rewrite pass (merge/useless, canonicalize, stabilize, specialize/fusion, backend-specific + in-place) and each Stage 2 backend in plain language with analogies, and move precise module/function references into collapsible "Under the hood" notes. Clarify how the backend (backend= / compile_kwargs / config.linker) and the sampler (nuts_sampler=: pymc/nutpie/numpyro/blackjax) are chosen, and make the PyTensor/sampler boundary explicit in the diagram. Co-authored-by: Cursor --- doc/compilation_for_pymc_users.md | 145 ++++++++++++++++++++++-------- 1 file changed, 110 insertions(+), 35 deletions(-) diff --git a/doc/compilation_for_pymc_users.md b/doc/compilation_for_pymc_users.md index 3626f5724e..4139f8f61c 100644 --- a/doc/compilation_for_pymc_users.md +++ b/doc/compilation_for_pymc_users.md @@ -137,54 +137,97 @@ PyMC hands PyTensor a symbolic graph — a recipe of mathematical operations, no ### Stage 1 — rewriting (the "optimizer") -PyTensor rewrites the graph to be faster and more numerically stable. This is where, for example, a naive `log(1 + x)` becomes a stable `log1p(x)`, duplicate sub-expressions get merged, and element-wise operations get fused together. Two things are worth internalising up front: **this stage is pure Python**, and its cost scales with the size of your graph (large hierarchical models have large graphs, so this can be a real chunk of "time before sampling"); and **this stage is backend-agnostic** — whether you end up on Numba, C, JAX, PyTorch, or MLX, almost all of the same rewriting happens first. +Think of the graph PyMC hands over as a **rough draft** of your model's math. Before running anything, PyTensor behaves like a careful editor working through that draft: it rewrites the math to run faster, to use less memory, and — especially important for probabilistic models — to be numerically safe. The key thing is that none of these edits change *what* the math computes. They only change *how* it's written down. The answer is identical; the path to it is better. -The subsections below walk through the boxes in the diagram's "Stage 1" lane, in order. Each names the module or function that implements it, so you can jump into the source. +Two things are worth holding onto before we look at the individual edits: -#### The optimizer pipeline (`optdb`) +- **It all happens in plain Python**, and the bigger your model, the longer it takes. This editing pass is a large part of the "why does sampling take a while to *start*?" pause. A big hierarchical model is a big draft, and a big draft takes longer to edit. +- **It happens the same way regardless of backend.** Whether you'll eventually run on Numba, C, JAX, PyTorch, or MLX, almost all of this rewriting happens first, in exactly the same way. -Every rewrite lives in one big, ordered registry: `pytensor.compile.mode.optdb`, an instance of `pytensor.graph.rewriting.db.SequenceDB`. Entries are registered at a numeric `position` (merge at `0`, `canonicalize` at `1`, `stabilize` at `1.5`, `specialize` at `2`, fusion at `49`, inplace at `49.5`+), and the pipeline runs them in that order. Each entry also carries *tags* such as `"fast_run"`, `"fast_compile"`, `"numba"`, or `"inplace"`. +The subsections below walk through each editing pass, in the order they run — matching the "Stage 1" lane in the diagram. Each one ends with a foldable **"Under the hood"** note that points at the exact PyTensor code, in case you ever want to go spelunking. You can safely skip those folds and still understand the whole pipeline. -Which entries actually run is decided by the {class}`pytensor.compile.mode.Mode` you compile with. A `Mode` pairs a *linker* (Stage 2) with an *optimizer query* (`pytensor.graph.rewriting.db.RewriteDatabaseQuery`) that selects rewrites by tag. For instance the default Numba mode queries `include=["fast_run", "numba"]`, while `FAST_COMPILE` queries `include=["fast_compile", "py_only"]` and so skips most of the pipeline. This is exactly the knob `FAST_COMPILE` and friends are turning. +#### The to-do list of edits -#### `merge` / `useless` +All the possible edits live in one long, **ordered to-do list**. Each edit has a number that fixes where it sits in the queue, and a few labels attached to it. When you compile, PyTensor does *not* run every edit on the list — it runs only the ones whose labels match the mode you picked. That is the real difference between `FAST_COMPILE` (run just a handful of quick edits, so you can start sampling sooner) and the default (run the full set, so sampling itself is fast). Choosing a compilation mode is really choosing *which edits are allowed to run*. -The first passes are cheap clean-ups. **Merge** (`pytensor.graph.rewriting.basic.MergeOptimizer`, registered as `merge1`, `merge1.1`, `merge2`, `merge3` at several positions) collapses identical sub-graphs so a repeated expression is computed once. **Useless** (the `"useless"` entry at position `0.6`, a `pytensor.graph.rewriting.db.TopoDB` over a `LocalGroupDB` of small `local_useless_*` node rewriters) strips nodes that don't affect the output, e.g. an identity reshape or a fill that broadcasts to a shape already in hand. +:::{dropdown} Under the hood +Every rewrite lives in `pytensor.compile.mode.optdb`, a `pytensor.graph.rewriting.db.SequenceDB`. Entries register at a numeric `position` (merge at `0`, canonicalize at `1`, stabilize at `1.5`, specialize at `2`, fusion at `49`, inplace at `49.5`+) and run in that order. Each carries *tags* like `"fast_run"`, `"fast_compile"`, `"numba"`, `"inplace"`. The {class}`pytensor.compile.mode.Mode` pairs a linker (Stage 2) with an *optimizer query* (`pytensor.graph.rewriting.db.RewriteDatabaseQuery`) that selects rewrites by tag: the default Numba mode queries `include=["fast_run", "numba"]`, while `FAST_COMPILE` queries `include=["fast_compile", "py_only"]`. +::: + +#### Tidying up: `merge` and `useless` + +The first edits are easy wins. **Merge** spots when the same calculation appears more than once and makes the model compute it a single time, then reuse the result — like noticing you've added up the same column of numbers twice and deciding to do it once. **Useless** deletes steps that can't possibly change the answer: adding `0`, multiplying by `1`, or reshaping an array into the shape it already has. Small savings each, but they also leave a cleaner draft for the heavier edits that follow. + +:::{dropdown} Under the hood +Merge is `pytensor.graph.rewriting.basic.MergeOptimizer` (registered as `merge1`, `merge1.1`, `merge2`, `merge3` at several positions). "Useless" is the `"useless"` entry at position `0.6`, a `pytensor.graph.rewriting.db.TopoDB` wrapping a `LocalGroupDB` of small `local_useless_*` node rewriters. +::: -#### `canonicalize` (fixed-point loop) +#### Speaking with one voice: `canonicalize` -**Canonicalize** (position `1`) rewrites the graph into a single normal form so that later passes only have to recognise one shape of each pattern. It is an `pytensor.graph.rewriting.db.EquilibriumDB`: it applies all of its member rewrites repeatedly until the graph stops changing (reaches *equilibrium* / a fixed point). Developers add a local rewrite here with the `@register_canonicalize` decorator from `pytensor.tensor.rewriting.basic`. Typical work: reordering and flattening `add`/`mul` chains, folding constants, and pushing operations into a canonical position. +There are usually many equally-valid ways to write the same expression. Left alone, that variety would force every later edit to recognise dozens of slightly different spellings of the same thing. **Canonicalize** prevents that by rewriting everything into one agreed-upon *standard form* — a house style for the math. Once the whole draft is written consistently, the later passes only have to look for one version of each pattern. This pass runs in a loop, applying its edits over and over until the draft stops changing (mathematicians call that reaching a *fixed point*); it keeps tidying until there's nothing left to tidy. -#### `stabilize` (numerically-stable forms) +:::{dropdown} Under the hood +Canonicalize is an `pytensor.graph.rewriting.db.EquilibriumDB` at position `1` — "equilibrium" being the run-until-nothing-changes loop. Developers add a rewrite to it with the `@register_canonicalize` decorator from `pytensor.tensor.rewriting.basic`. Typical work: flattening `add`/`mul` chains, folding constants, and moving operations into a canonical position. +::: -**Stabilize** (position `1.5`, also an `EquilibriumDB`, populated via `@register_stabilize`) swaps numerically fragile expressions for equivalent stable ones — the step PyMC users feel most, because log-probabilities are full of logs and exps. Concrete examples from `pytensor.tensor.rewriting.math`: `local_log1p` turns `log(1 + x)` into `log1p(x)`, `local_expm1` turns `exp(x) - 1` into `expm1(x)`, and `local_log1p_plusminus_exp` turns `log1p(exp(x))` into `log1pexp(x)` (the softplus). These avoid catastrophic cancellation and overflow that the naive forms suffer near the extremes. +#### Keeping the numbers safe: `stabilize` -#### `specialize` + `fusion` +This is the pass PyMC users benefit from the most, even though they almost never see it work. Here is the problem it solves. Computers store numbers with only so many digits of precision, and some formulas that are perfectly correct on paper fall apart on a real computer. -**Specialize** (position `2`, `EquilibriumDB`, `@register_specialize`) replaces general operations with faster special-cased ones (for example rewriting `x ** 2` or `sum(x**2)` into cheaper equivalents). **Fusion** then merges many element-wise operations into a single `Composite` op so the runtime walks the data once instead of allocating a temporary per operation: see the `"elemwise_fusion"` sequence (the `FusionOptimizer`) and `"add_mul_fusion"` in `pytensor.tensor.rewriting.elemwise`. Fusion is tagged `"fusion"`, which is why `FAST_COMPILE` (and JAX, which does its own fusion) leave it out. +Take `log(1 + x)` when `x` is tiny — say `0.000000001`. The computer first works out `1 + x`, but it can't keep all those digits, so it rounds the result to exactly `1`. Then `log(1) = 0`, and your small-but-real value has simply vanished. The mathematically-identical function `log1p(x)` computes the same thing *without ever forming `1 + x`*, so the precision survives. **Stabilize** automatically swaps these fragile formulas for safe equivalents like that one. -#### Backend-specific rewrites + inplace / destroy handler +:::{tip} +If you've ever wondered why PyMC so rarely spits out `NaN` or `inf` from ordinary-looking expressions full of logs and exponentials, this pass is a big part of the answer. It is quietly rewriting the danger out of your log-probability. +::: -Two kinds of rewrites run near the end. **Backend-specific** rewrites are pulled in by the mode's tag query — the Numba mode adds `"numba"`-tagged rewrites (`pytensor.tensor.rewriting.numba`), the JAX mode adds `"jax"`-tagged ones (`pytensor.tensor.rewriting.jax`), and so on; conversely C-only ops and BLAS rewrites are *excluded* for the array-library backends. **Inplace** rewrites (tagged `"inplace"`) let ops write into their inputs' memory to avoid allocations: `AddDestroyHandler` (position `49.5`) first attaches a `pytensor.graph.destroyhandler.DestroyHandler` that tracks which buffers may be clobbered, then rewrites like `InplaceElemwiseOptimizer` (`pytensor.tensor.rewriting.elemwise`, position `50.5`) switch ops to their in-place variants. JAX excludes inplace entirely, since it manages memory itself. +:::{dropdown} Under the hood +Stabilize is an `EquilibriumDB` at position `1.5`, populated via `@register_stabilize`. Concrete examples from `pytensor.tensor.rewriting.math`: `local_log1p` (`log(1 + x)` → `log1p(x)`), `local_expm1` (`exp(x) - 1` → `expm1(x)`), and `local_log1p_plusminus_exp` (`log1p(exp(x))` → `log1pexp(x)`, the softplus). These avoid the cancellation and overflow the naive forms suffer at the extremes. +::: -#### Rewritten FunctionGraph +#### Shortcuts and fewer passes over the data: `specialize` and `fusion` -The output of Stage 1 is a new {class}`pytensor.graph.fg.FunctionGraph` — the same inputs and outputs, but a leaner, more stable, backend-tailored computation. Nothing has been compiled yet; that is Stage 2's job. (To inspect this graph without compiling a full function, apply the rewrites directly with `pytensor.graph.rewriting.utils.rewrite_graph(y, include=("canonicalize", "specialize"))` and print the result with `pytensor.dprint`.) +Two more speed-ups, both about doing less work. **Specialize** swaps a general-purpose operation for a faster special-case version whenever one exists — the way you'd compute `x²` as `x * x` instead of starting up a general "raise to any power" routine. + +**Fusion** fixes a subtler waste. Imagine your model computes `a * b + c` element by element across big arrays. Done naively, the computer builds one whole temporary array to hold `a * b`, then another to hold `+ c` — a lot of memory shuffled around for nothing. Fusion merges these element-wise steps so each element gets *all* of its operations in a single pass, with no temporaries in between. Picture an assembly line where each item is fully finished at one station, instead of the entire batch being carried back and forth between stations. + +:::{dropdown} Under the hood +Specialize is an `EquilibriumDB` at position `2` (`@register_specialize`). Fusion is the `"elemwise_fusion"` sequence — the `FusionOptimizer`, which builds a single `Composite` op — plus `"add_mul_fusion"`, both in `pytensor.tensor.rewriting.elemwise`. Fusion is tagged `"fusion"`, which is why `FAST_COMPILE` (and JAX, which fuses on its own) leave it out. +::: + +#### Last-minute, backend-aware edits: backend-specific rewrites and in-place ops + +The final edits depend on *where* your model is headed and on squeezing out memory. **Backend-specific** edits are small tweaks tailored to the backend you chose — a few extra tricks that only make sense for Numba, or only for JAX, and so on. Tricks that only help the C backend are dropped when you're not using it. + +**In-place** edits save memory by letting an operation write its result directly over its input's memory, instead of allocating a brand-new array — like reusing the same sheet of scratch paper rather than grabbing a fresh one every time. Overwriting is risky, though: what if some other part of the graph still needs that data? So before any in-place edit runs, a "destroy handler" acts as a safety inspector, tracking exactly which chunks of memory are safe to reuse. (JAX skips in-place edits altogether — it manages memory in its own way.) + +:::{dropdown} Under the hood +Backend-specific rewrites are pulled in by the mode's tag query (`pytensor.tensor.rewriting.numba`, `pytensor.tensor.rewriting.jax`, …); C-only and BLAS rewrites are excluded for the array-library backends. In-place rewrites are tagged `"inplace"`: `AddDestroyHandler` (position `49.5`) attaches a `pytensor.graph.destroyhandler.DestroyHandler` that tracks clobberable buffers, then rewrites like `InplaceElemwiseOptimizer` (`pytensor.tensor.rewriting.elemwise`, position `50.5`) switch ops to their in-place variants. +::: + +#### The finished draft + +The result of all this editing is a new graph: the same inputs and outputs as before, but leaner, numerically safer, and tailored to your backend. Nothing has actually *run* yet — we've only improved the recipe. Turning that recipe into a program the computer can execute is Stage 2's job. + +:::{dropdown} Under the hood +The output is a rewritten {class}`pytensor.graph.fg.FunctionGraph`. To see it without compiling a full function, apply the rewrites directly with `pytensor.graph.rewriting.utils.rewrite_graph(y, include=("canonicalize", "specialize"))` and print the result with `pytensor.dprint`. +::: ### Stage 2 — linking (choosing a backend) -The rewritten graph is turned into an executable. *How* depends on the linker/backend: +After Stage 1 you have a polished recipe — but a recipe is just instructions, not a meal. **Linking** is the step that turns those instructions into a real program the computer can run quickly. The recipe is still written in PyTensor's own internal language; linking translates it into a language that actually executes, and (for most backends) compiles it. *How* that translation happens, and how fast the result runs, depends on the **backend** you chose. + +Here's the lay of the land before we dig in: -| Backend | How it compiles | Needs a system compiler? | Notes | +| Backend | How it turns the recipe into a program | Needs a system compiler? | Good to know | |---|---|---|---| -| **Numba** *(default)* | Converts the graph to Python, then JIT-compiles via LLVM | No (ships as wheels) | This is why `pip install pymc` "just works" without conda. First compile is slow; results are cached on disk. | -| **C / CVM** | Generates C++ source and compiles `.so` files with `g++`/`clang++` | Yes | Historically the default. Strong on-disk module cache. Now opt-in. | -| **JAX** | Converts to JAX, compiles via `jax.jit` → XLA | No | Great for GPU/TPU and `vmap`-style workflows. | -| **PyTorch** | Converts to PyTorch, compiles via `torch.compile` (TorchDynamo/Inductor) | No | Runs on PyTorch's CPU/GPU devices. | -| **MLX** | Converts to MLX, compiles via `mx.compile` | No | Targets Apple Silicon GPUs. | -| **Python VM** | Runs each operation's pure-Python implementation | No | Slowest at runtime, fastest to "compile". Handy for debugging. | +| **Numba** *(default)* | Converts the recipe to Python, then compiles it to fast machine code the first time you run it (via LLVM) | No (ships ready-to-use) | Why `pip install pymc` "just works" without conda. First run is slow; results are cached on disk. | +| **C / CVM** | Writes out C++ source and compiles it with `g++`/`clang++` | Yes | Historically the default; very strong on-disk cache. Now opt-in. | +| **JAX** | Converts to JAX, compiles via `jax.jit` → XLA | No | The path to GPU/TPU. | +| **PyTorch** | Converts to PyTorch, compiles via `torch.compile` | No | Runs on PyTorch's CPU/GPU devices. | +| **MLX** | Converts to MLX, compiles via `mx.compile` | No | Targets the GPU in Apple Silicon Macs. | +| **Python VM** | Just runs each operation in plain Python — no compiling | No | Slowest to run, instant to "build". Great for debugging. | -The default backend resolves from `config.linker = "auto"`, which currently maps to **Numba**. +The default backend is **Numba** (it's what `config.linker = "auto"` resolves to). #### How the backend gets chosen @@ -220,33 +263,65 @@ f = pytensor.function([x], y, mode="NUMBA") The two choices — backend and sampler — are made at different points in the pipeline (see the two blue decision diamonds in the diagram), and as noted below they're mostly independent. The exception is that JAX-based samplers force the JAX backend. -#### How linking works (the shared mechanism) +#### The one idea behind (almost) every backend + +Most backends work the same way under the hood, and once you see it the rest is just details. PyTensor keeps a kind of **phrasebook**: for every operation in your graph, it knows how to say that operation in the target language. To link a whole graph, it translates each operation in turn, stitches the translations together into one program, and then lets the backend compile that program. Adding support for a new operation in a backend is just adding one new phrase to that backend's phrasebook. -Every backend is implemented as a *linker*, a subclass of `pytensor.link.basic.Linker`; the {class}`pytensor.compile.mode.Mode`'s linker is what Stage 2 runs. Most linkers share one trick: a single-dispatch registry that translates each PyTensor `Op` into the target language. Numba uses `pytensor.link.numba.dispatch.numba_funcify`, JAX uses `pytensor.link.jax.dispatch.jax_funcify`, PyTorch uses `pytensor.link.pytorch.dispatch.pytorch_funcify`, and MLX uses `pytensor.link.mlx.dispatch.mlx_funcify`. To add backend support for a new `Op`, you register a function with the matching dispatcher. The C and Python backends are the exception: rather than a `*_funcify` registry, they ask each `Op` directly for code through `Op.c_code` or run its `Op.perform` method. The subsections below cover each branch of the diagram's "Stage 2" lane. +The C and Python backends are the two exceptions — instead of a phrasebook they ask each operation directly for its hand-written C code, or simply run the operation's plain-Python implementation. The rest of this section takes the backends one at a time (the branches of the diagram's "Stage 2" lane). + +:::{dropdown} Under the hood +Each backend is a *linker* (a `pytensor.link.basic.Linker` subclass); the {class}`pytensor.compile.mode.Mode`'s linker is what Stage 2 runs. The "phrasebook" is a single-dispatch registry: `pytensor.link.numba.dispatch.numba_funcify`, `pytensor.link.jax.dispatch.jax_funcify`, `pytensor.link.pytorch.dispatch.pytorch_funcify`, `pytensor.link.mlx.dispatch.mlx_funcify`. The C and Python backends instead use `Op.c_code` and `Op.perform`. +::: -#### Numba (default) +#### Numba (the default) -`pytensor.link.numba.linker.NumbaLinker`. `numba_funcify` walks the rewritten `FunctionGraph` and emits one Python function that wires together each op's Numba implementation; `pytensor.link.numba.dispatch.numba_njit` then JIT-compiles it with [`numba.njit`](https://numba.readthedocs.io/en/stable/user/jit.html) via LLVM. The first call pays the compilation cost; afterwards it runs as native machine code. With `config.numba__cache = True` (the default) the compiled result is written to an on-disk cache (keyed per op, see `numba_funcify_and_cache_key`), so a later session can skip recompiling. Numba ships as wheels, which is why `pip install pymc` "just works" without a system compiler. +Numba turns your recipe into Python code, then hands it to a **just-in-time (JIT) compiler** that converts it into fast machine code the very first time you run it. That's why the first call is slow — it includes the compiling — while every call afterwards is quick. Better still, the compiled result is saved to disk, so even a brand-new Python session can usually skip the wait. + +Here's the intuition: the first time you cook an unfamiliar dish you're slow, reading each line and hunting for ingredients. By the tenth time it's second nature. Numba writes down that "second nature" version so it never has to relearn it. And because Numba comes ready-to-use when you `pip install`, there's no compiler for you to set up — which is exactly why `pip install pymc` "just works". + +:::{dropdown} Under the hood +`pytensor.link.numba.linker.NumbaLinker`. `numba_funcify` emits one Python function wiring together each op's Numba implementation; `pytensor.link.numba.dispatch.numba_njit` JIT-compiles it with [`numba.njit`](https://numba.readthedocs.io/en/stable/user/jit.html) via LLVM. With `config.numba__cache = True` (default) the result is cached on disk (see `numba_funcify_and_cache_key`). +::: #### C / CVM -`pytensor.link.c.basic.CLinker` generates C++ for each op from its `Op.c_code`, compiles a `.so` with the system `g++`/`clang++`, and caches the module in your compile directory. The default `"cvm"` linker is `pytensor.link.vm.VMLinker` with `use_cloop=True`: a hybrid where a fast C "virtual machine" loop steps through op *thunks*, each of which may be C-compiled or fall back to Python. Pure `"c"` (`CLinker` / `OpWiseCLinker`) compiles as much of the graph into C as possible. This family was historically the default and has a strong on-disk module cache, but it needs a working C++ toolchain — hence the move to Numba for the pip story. +The C backend writes out C++ source code and compiles it using a real C++ compiler (`g++` or `clang++`) installed on your machine. It runs fast and keeps an excellent on-disk cache — but it only works if those build tools are present. That requirement is the whole reason it's no longer the default: plenty of people who `pip install pymc` don't have a C++ compiler set up. The everyday `"cvm"` variant is a practical hybrid — a small, fast loop written in C marches through your operations, each of which is either compiled C or falls back to Python. + +:::{dropdown} Under the hood +`pytensor.link.c.basic.CLinker` generates C++ from each op's `Op.c_code`, compiles a `.so`, and caches it in your compile directory. `"cvm"` is `pytensor.link.vm.VMLinker` with `use_cloop=True` (the C loop steps through op *thunks*); pure `"c"` is `CLinker` / `OpWiseCLinker`. +::: #### JAX -`pytensor.link.jax.linker.JAXLinker`. `jax_funcify` translates the graph into a JAX-traceable Python function, which the linker hands to {func}`jax.jit` → XLA. This is the GPU/TPU-friendly path, and it's also the backend the JAX samplers (`numpyro`, `blackjax`, and `nutpie` in JAX mode) consume directly. +JAX translates the recipe into JAX code and lets JAX's own compiler (XLA) take over. This is the route to **GPUs and TPUs**, and it's the backend the JAX-based samplers (`numpyro`, `blackjax`, and `nutpie` in its JAX mode) plug straight into. + +:::{dropdown} Under the hood +`pytensor.link.jax.linker.JAXLinker`. `jax_funcify` builds a JAX-traceable function that the linker hands to {func}`jax.jit` → XLA. +::: #### PyTorch -`pytensor.link.pytorch.linker.PytorchLinker`. `pytorch_funcify` builds a PyTorch function, compiled with {func}`torch.compile` (TorchDynamo + Inductor), running on PyTorch's CPU/GPU devices. +Same idea, using PyTorch instead. The recipe is translated into PyTorch and compiled with {func}`torch.compile`, running on whatever CPU or GPU device PyTorch is set up to use. Handy when the rest of your stack already lives in PyTorch. + +:::{dropdown} Under the hood +`pytensor.link.pytorch.linker.PytorchLinker`; `pytorch_funcify` builds the function, compiled via {func}`torch.compile` (TorchDynamo + Inductor). +::: #### MLX -`pytensor.link.mlx.linker.MLXLinker`. `mlx_funcify` builds an MLX function compiled lazily with `mx.compile`, targeting Apple Silicon GPUs. +The same idea once more, this time aimed at the GPU built into Apple Silicon Macs. + +:::{dropdown} Under the hood +`pytensor.link.mlx.linker.MLXLinker`; `mlx_funcify` builds an MLX function compiled lazily with `mx.compile`. +::: #### Python VM -The reference backend: `pytensor.link.basic.PerformLinker` (or `pytensor.link.vm.VMLinker` with `use_cloop=False`) steps through the graph calling each op's `Op.perform` in pure Python. There is essentially no compilation step, so it is the fastest to "build" and the slowest to run — ideal for debugging, and what `FAST_COMPILE` uses. +The simplest backend of all: **no compiling whatsoever**. PyTensor just runs each operation one at a time in plain Python. That makes it the slowest to actually run, but instant to "build" and by far the easiest to debug (you can drop into any step with ordinary Python tools). It's what `FAST_COMPILE` uses, and it's a great companion while you're still iterating on a model's structure. + +:::{dropdown} Under the hood +`pytensor.link.basic.PerformLinker` (or `pytensor.link.vm.VMLinker` with `use_cloop=False`) steps through the graph calling each op's `Op.perform`. +::: ### Runtime