From a0f228bc045bdfbdf7c47603a8aa7207d82c5fed Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Sat, 13 Jun 2026 22:45:30 +0100 Subject: [PATCH 1/7] docs: add "Understanding PyTensor compilation" core notebook page Add an explanation-type page to the core notebooks that describes how PyMC hands a model's math to PyTensor and what PyTensor does with it: the two-stage compilation pipeline (graph rewriting then backend linking), the available backends (Numba/C/JAX/PyTorch/MLX/VM), and the boundary between PyTensor compilation and the MCMC sampler engine (pymc/nutpie/numpyro/blackjax). Complements the existing "PyMC and PyTensor" notebook (which covers how models become graphs). Includes a pipeline diagram (committed as SVG, with the Mermaid source alongside) and links it into the core-notebooks toctree. Co-authored-by: Cursor --- docs/source/images/pytensor_compilation.mmd | 91 ++++++ docs/source/images/pytensor_compilation.svg | 1 + docs/source/learn/core_notebooks/index.md | 1 + .../core_notebooks/pytensor_compilation.md | 273 ++++++++++++++++++ 4 files changed, 366 insertions(+) create mode 100644 docs/source/images/pytensor_compilation.mmd create mode 100644 docs/source/images/pytensor_compilation.svg create mode 100644 docs/source/learn/core_notebooks/pytensor_compilation.md diff --git a/docs/source/images/pytensor_compilation.mmd b/docs/source/images/pytensor_compilation.mmd new file mode 100644 index 0000000000..2ee56e394a --- /dev/null +++ b/docs/source/images/pytensor_compilation.mmd @@ -0,0 +1,91 @@ +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 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)"] + 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 = 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: 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; diff --git a/docs/source/images/pytensor_compilation.svg b/docs/source/images/pytensor_compilation.svg new file mode 100644 index 0000000000..169f74a3c5 --- /dev/null +++ b/docs/source/images/pytensor_compilation.svg @@ -0,0 +1 @@ +

Runtime — MCMC sampling (the sampler engine, NOT PyTensor)

STAGE 2 — Linking (backend-specific)

STAGE 1 — Graph rewriting (pure Python, backend-agnostic)

Graph preparation (pure Python)

PyMC layer

Python VM backend

MLX backend (Apple Silicon GPU)

PyTorch backend

JAX backend

C / CVM backend (needs C++ compiler)

Numba backend (default, pip-friendly)

'pymc'

'nutpie' (default if installed)

'numpyro'

'blackjax'

'numba' (default)

'c' (→ cvm)

'jax'

'pytorch'

'mlx' (Apple GPU)

'py' / 'vm'

PyMC model
(pm.Model: priors, likelihood)

Build symbolic graph
logp / dlogp / random fns

pytensor.function(inputs, outputs, mode=...)

clone graph, apply givens/updates,
collect shared variables

wrap in FunctionGraph

run optimizer pipeline (optdb)

merge / useless

canonicalize
(equilibrium, fixed-point loop)

stabilize
numerically-stable forms

specialize + fusion

backend-specific rewrites
+ inplace / destroy handler

Rewritten FunctionGraph

Which backend (linker)?
pm.sample(backend=...)
or compile_kwargs mode=... / config.linker

graph -> Python

numba.njit JIT via LLVM

on-disk cache
(numba__cache=True)

generate C++ source

g++/clang++ subprocess

cached .so modules

graph -> JAX

jax.jit -> XLA

graph -> PyTorch

torch.compile
(TorchDynamo/Inductor)

graph -> MLX

mx.compile (lazy JIT)

Op.perform in pure Python

Compiled callable — PyTensor's output
θ (a point in parameter space) → logp(θ), ∇logp(θ)
(observed data baked in as constants)

Sampler loop:
propose θ → call the compiled logp/∇logp → accept/reject → repeat (many times)

Which sampler engine?
pm.sample(nuts_sampler=...)

PyMC NUTS
pure-Python loop
(calls a callable from any CPU backend, e.g. C / Numba)

nutpie / nuts-rs
Rust loop
(needs a Numba or JAX callable)

NumPyro NUTS
JAX loop
(needs a JAX callable)

BlackJAX NUTS
JAX loop
(needs a JAX callable)

 diff --git a/docs/source/learn/core_notebooks/index.md b/docs/source/learn/core_notebooks/index.md index 4e7d7353b1..de9782e44f 100644 --- a/docs/source/learn/core_notebooks/index.md +++ b/docs/source/learn/core_notebooks/index.md @@ -9,6 +9,7 @@ model_comparison posterior_predictive dimensionality pymc_pytensor +pytensor_compilation dims_module GLM_linear Gaussian_Processes diff --git a/docs/source/learn/core_notebooks/pytensor_compilation.md b/docs/source/learn/core_notebooks/pytensor_compilation.md new file mode 100644 index 0000000000..ce95c7eb6c --- /dev/null +++ b/docs/source/learn/core_notebooks/pytensor_compilation.md @@ -0,0 +1,273 @@ +(pytensor_compilation)= +# Understanding PyTensor compilation + +:::{seealso} +This page is a companion to {ref}`pymc_pytensor`. That notebook explains *how* a PyMC model becomes a PyTensor graph; this one picks up where it leaves off and explains what PyTensor then *does* with that graph to make sampling fast. +::: + +:::{note} +This page is deliberately *opinionated* and slightly broader in scope than a pure PyTensor reference: 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. +::: + +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 + +:::{figure} ../../images/pytensor_compilation.svg +:width: 100% +:alt: A flowchart of the PyTensor compilation pipeline, from the PyMC model through graph preparation, Stage 1 graph rewriting, backend selection, and Stage 2 linking, to the compiled callable consumed by an MCMC sampler engine. + +The full pipeline. Everything from the PyMC model down to the **compiled callable** is PyTensor's job; everything below it is the *sampler engine*. The two blue diamonds are the two choices you control: the **backend** (how the math is compiled) and the **sampler engine** (what repeatedly evaluates it). +::: + +## 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 {class}`~pytensor.graph.fg.FunctionGraph`. This is bookkeeping; it's cheap. + +### Stage 1 — rewriting (the "optimizer") + +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. + +Two things are worth holding onto before we look at the individual edits: + +- **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. + +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. + +#### The to-do list of edits + +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*. + +:::{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. +::: + +#### Speaking with one voice: `canonicalize` + +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. + +:::{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. +::: + +#### Keeping the numbers safe: `stabilize` + +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. + +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. + +:::{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. +::: + +:::{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. +::: + +#### Shortcuts and fewer passes over the data: `specialize` and `fusion` + +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) + +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 turns the recipe into a program | Needs a system compiler? | Good to know | +|---|---|---|---| +| **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 is **Numba** (it's what `config.linker = "auto"` resolves to). + +#### 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. + +#### 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. + +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 (the default) + +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 + +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 + +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 [`jax.jit`](https://docs.jax.dev/en/latest/_autosummary/jax.jit.html) → XLA. +::: + +#### PyTorch + +Same idea, using PyTorch instead. The recipe is translated into PyTorch and compiled with [`torch.compile`](https://docs.pytorch.org/docs/stable/generated/torch.compile.html), 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 `torch.compile` (TorchDynamo + Inductor). +::: + +#### MLX + +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 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 + +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 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. + +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 + +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}`pymc_pytensor` — how PyMC models become PyTensor graphs in the first place. +- {ref}`graph rewrites catalogue ` — the rewrites applied in Stage 1. +- {ref}`modes and linkers ` — modes and linkers in more depth. +- {doc}`Mode / linker API reference `. From 4b25f25c4fa9ad84d8516a76b4880999f5d7d727 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Sat, 20 Jun 2026 08:41:54 +0100 Subject: [PATCH 2/7] actual human edits --- docs/source/learn/core_notebooks/index.md | 2 +- .../core_notebooks/pytensor_compilation.md | 52 +++++++++---------- 2 files changed, 27 insertions(+), 27 deletions(-) diff --git a/docs/source/learn/core_notebooks/index.md b/docs/source/learn/core_notebooks/index.md index de9782e44f..8bc4af59bd 100644 --- a/docs/source/learn/core_notebooks/index.md +++ b/docs/source/learn/core_notebooks/index.md @@ -8,8 +8,8 @@ pymc_overview model_comparison posterior_predictive dimensionality -pymc_pytensor pytensor_compilation +pymc_pytensor dims_module GLM_linear Gaussian_Processes diff --git a/docs/source/learn/core_notebooks/pytensor_compilation.md b/docs/source/learn/core_notebooks/pytensor_compilation.md index ce95c7eb6c..b417603352 100644 --- a/docs/source/learn/core_notebooks/pytensor_compilation.md +++ b/docs/source/learn/core_notebooks/pytensor_compilation.md @@ -1,24 +1,26 @@ (pytensor_compilation)= -# Understanding PyTensor compilation +# PyTensor, WTF? -:::{seealso} -This page is a companion to {ref}`pymc_pytensor`. That notebook explains *how* a PyMC model becomes a PyTensor graph; this one picks up where it leaves off and explains what PyTensor then *does* with that graph to make sampling fast. -::: - -:::{note} -This page is deliberately *opinionated* and slightly broader in scope than a pure PyTensor reference: 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. -::: +**Author:** [Benjamin Vincent](https://github.com/drbenvincent) -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: +In your day-to-day modeling work, it is possible to write models exclusively in PyMC and never know about, or write any, PyTensor code. But PyTensor is a _crucial_ part of the modeling stack, and does a lot of work on your behalf every time you sample. Understanding what it is and how it works can help 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.) +- **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**. +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. + +$$ +\theta \to \big(\log(p(\theta)), \nabla \log(p(\theta))\big) +$$ + +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 @@ -44,7 +46,7 @@ Two things are worth holding onto before we look at the individual edits: - **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. -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. +The subsections below walk through each editing pass, in the order they run — matching the "Stage 1" lane in the diagram. For a full catalogue of every available rewrite, see the {ref}`graph rewrites reference `. #### The to-do list of edits @@ -120,14 +122,14 @@ Here's the lay of the land before we dig in: | Backend | How it turns the recipe into a program | Needs a system compiler? | Good to know | |---|---|---|---| -| **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. | +| **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" (from PyMC 6.0 onwards) 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 (pre PyMC 6.0); 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 is **Numba** (it's what `config.linker = "auto"` resolves to). +The default backend is **Numba** (it's what `config.linker = "auto"` resolves to). For a deeper treatment of how modes and linkers work, see {ref}`modes and linkers `. #### How the backend gets chosen @@ -161,7 +163,14 @@ pytensor.config.linker = "jax" 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. +The two choices — backend and sampler — are made at different points in the pipeline (see the two blue decision diamonds in the diagram), and they're mostly independent. You pick the engine with `pm.sample(nuts_sampler=...)`: + +- **`"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. + +The one coupling to be aware of: 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 one idea behind (almost) every backend @@ -231,15 +240,6 @@ An MCMC sampler calls this function many times, proposing new parameter points a **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. - 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 6ae42461f60d68d35e72e1eb672db6b30f3b0c29 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Sat, 20 Jun 2026 08:46:01 +0100 Subject: [PATCH 3/7] run pre-commit to fix things --- docs/source/learn/core_notebooks/pytensor_compilation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/learn/core_notebooks/pytensor_compilation.md b/docs/source/learn/core_notebooks/pytensor_compilation.md index b417603352..9e01f18aef 100644 --- a/docs/source/learn/core_notebooks/pytensor_compilation.md +++ b/docs/source/learn/core_notebooks/pytensor_compilation.md @@ -12,7 +12,7 @@ This guide gives you a mental model of the pipeline and the choices available, w ## 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. +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. $$ \theta \to \big(\log(p(\theta)), \nabla \log(p(\theta))\big) From d853ee8b740a909526a067b63dbc0821849b7fac Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Fri, 26 Jun 2026 18:48:42 +0100 Subject: [PATCH 4/7] docs: trim PyTensor compilation guide to overview level Reposition the page as a high-level overview per review feedback: broaden the mental model beyond logp+gradient, remove all Under the hood dropdowns, fix factual errors (useless/canonicalize, pm.Data), and soften the Stage 1 caching claim. Co-authored-by: Cursor --- .../core_notebooks/pytensor_compilation.md | 103 +++++------------- 1 file changed, 27 insertions(+), 76 deletions(-) diff --git a/docs/source/learn/core_notebooks/pytensor_compilation.md b/docs/source/learn/core_notebooks/pytensor_compilation.md index 9e01f18aef..b419a0ea39 100644 --- a/docs/source/learn/core_notebooks/pytensor_compilation.md +++ b/docs/source/learn/core_notebooks/pytensor_compilation.md @@ -12,11 +12,11 @@ This guide gives you a mental model of the pipeline and the choices available, w ## 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. +When PyMC fits or samples a model, it asks PyTensor to turn symbolic math into fast, callable functions. What gets compiled depends on what you're doing: -$$ -\theta \to \big(\log(p(\theta)), \nabla \log(p(\theta))\big) -$$ +- **Gradient-based MCMC** (NUTS, HMC): log-probability and its gradient — $\theta \to (\log p(\theta), \nabla \log p(\theta))$ +- **Gradient-free MCMC** (Metropolis, Slice, DEMetropolis): log-probability only — $\theta \to \log p(\theta)$ +- **Forward sampling** (`sample_posterior_predictive`, `sample_prior_predictive`): draws from the generative model — no log-probability at all 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. @@ -35,16 +35,17 @@ The full pipeline. Everything from the PyMC model down to the **compiled callabl ### 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 {class}`~pytensor.graph.fg.FunctionGraph`. This is bookkeeping; it's cheap. +PyMC hands PyTensor a symbolic graph — a recipe of mathematical operations, not yet any numbers. PyTensor clones it and prepares it for rewriting. This is bookkeeping; it's cheap. ### Stage 1 — rewriting (the "optimizer") 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. -Two things are worth holding onto before we look at the individual edits: +Three things are worth holding onto before we look at the individual edits: - **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. +- **Vectorized code helps here.** PyTensor rewrites the *graph*, not Python loops. NumPy-style vectorized expressions produce a compact graph the rewriter can simplify and translate cleanly; a Python `for` loop builds a much larger, deeper graph — even though the backend that eventually runs is not Python. The subsections below walk through each editing pass, in the order they run — matching the "Stage 1" lane in the diagram. For a full catalogue of every available rewrite, see the {ref}`graph rewrites reference `. @@ -52,25 +53,13 @@ The subsections below walk through each editing pass, in the order they run — 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*. -:::{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` +#### Tidying up: `merge` -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. -::: +The first edit is an easy win. **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. #### Speaking with one voice: `canonicalize` -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. - -:::{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. -::: +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. Typical tidy-ups include flattening chains of additions and multiplications, folding constants, and removing terms that cannot change the answer (adding `0`, multiplying by `1`). 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 (reaching an *equilibrium*); it keeps tidying until there's nothing left to tidy. #### Keeping the numbers safe: `stabilize` @@ -82,42 +71,28 @@ Take `log(1 + x)` when `x` is tiny — say `0.000000001`. The computer first wor 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. ::: -:::{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. -::: - #### Shortcuts and fewer passes over the data: `specialize` and `fusion` 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) 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. +Most backends translate the graph operation by operation into a target language, stitch the pieces into one program, and let the backend compile it. The C and Python backends are simpler: C generates source code per operation; Python just runs each operation in sequence. + Here's the lay of the land before we dig in: | Backend | How it turns the recipe into a program | Needs a system compiler? | Good to know | @@ -163,7 +138,7 @@ pytensor.config.linker = "jax" 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 they're mostly independent. You pick the engine with `pm.sample(nuts_sampler=...)`: +The two choices — backend and sampler — are made at different points in the pipeline (see the two blue decision diamonds in the diagram), and they're mostly independent. For NUTS specifically, you pick the engine with `pm.sample(nuts_sampler=...)`: - **`"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. @@ -172,75 +147,51 @@ The two choices — backend and sampler — are made at different points in the The one coupling to be aware of: 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 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. - -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 (the default) 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 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 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 [`jax.jit`](https://docs.jax.dev/en/latest/_autosummary/jax.jit.html) → XLA. -::: - #### PyTorch Same idea, using PyTorch instead. The recipe is translated into PyTorch and compiled with [`torch.compile`](https://docs.pytorch.org/docs/stable/generated/torch.compile.html), 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 `torch.compile` (TorchDynamo + Inductor). -::: - #### MLX 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 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 -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.) +What the compiled callable does depends on the task: + +| Task | What PyTensor compiles | Typical samplers | +|---|---|---| +| Gradient MCMC | $\theta \to (\log p(\theta), \nabla \log p(\theta))$ | NUTS (pymc, nutpie, numpyro, blackjax), HMC | +| Gradient-free MCMC | $\theta \to \log p(\theta)$ | Metropolis, Slice, DEMetropolis | +| Forward sampling | draws from the generative model | `sample_posterior_predictive`, `sample_prior_predictive` | + +PyMC may compile several such functions for one model — log-density and gradient for sampling, plus separate functions for prior or posterior predictive draws. + +Data can be baked in or left as a runtime input. Observed values passed as plain NumPy arrays are folded into the compiled function as constants. Values stored in `pm.Data` containers stay mutable: updating them with `pm.set_data` does not require recompilation. If you want data and dimension lengths treated as compile-time constants (for stronger rewrites), use the `freeze_dims_and_data` model transform. -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. +An MCMC sampler calls its compiled function many times, proposing new parameter points and reading back log-probability (and gradient, when needed) 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. -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. +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 function the Rust loop evaluates. ## Caching: why the second run is faster @@ -249,7 +200,7 @@ PyTensor caches the **linked artifact** so repeated builds can skip recompilatio - **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. +Note a current limitation that matters in practice: today, PyTensor's on-disk cache covers the *linked* output (Stage 2), but Stage 1 graph rewriting is re-run on every new compilation of a model. PyMC is improving this — compiled model functions are increasingly cached at the model level so repeat work on structurally identical graphs can be skipped. ## Practical tips for PyMC users From 86fca9094a602f14a8473cf2d39e3592b3eda92c Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Fri, 26 Jun 2026 18:57:43 +0100 Subject: [PATCH 5/7] docs: clarify pm.Data mutability and recompilation trade-off Expand the data handling paragraph in the Runtime section to explain why pm.Data avoids recompilation (runtime input vs compile-time constant) and when the trade-off matters. Co-authored-by: Cursor --- docs/source/learn/core_notebooks/pytensor_compilation.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/source/learn/core_notebooks/pytensor_compilation.md b/docs/source/learn/core_notebooks/pytensor_compilation.md index b419a0ea39..6409505c7a 100644 --- a/docs/source/learn/core_notebooks/pytensor_compilation.md +++ b/docs/source/learn/core_notebooks/pytensor_compilation.md @@ -185,7 +185,14 @@ What the compiled callable does depends on the task: PyMC may compile several such functions for one model — log-density and gradient for sampling, plus separate functions for prior or posterior predictive draws. -Data can be baked in or left as a runtime input. Observed values passed as plain NumPy arrays are folded into the compiled function as constants. Values stored in `pm.Data` containers stay mutable: updating them with `pm.set_data` does not require recompilation. If you want data and dimension lengths treated as compile-time constants (for stronger rewrites), use the `freeze_dims_and_data` model transform. +The way observed data is handled depends on how it was declared: + +- **Plain NumPy arrays** are folded into the compiled function as constants. The compiler can see those values at build time and use them to simplify the graph — but if you want to update the data you have to recompile from scratch. +- **`pm.Data` containers** are kept as *runtime inputs* that the function reads each time it is called, not at build time. The function is compiled once; you can then call `pm.set_data(...)` to swap the values in and the compiled function simply reads the new data on its next call — no recompilation needed. The trade-off is that the compiler cannot fold those values in to simplify the graph. + +This matters most in iterative workflows — posterior predictive checks over many data splits, or updating a model's observations between sampling runs — where recompiling each time would dominate wall time. + +If you want data and dimension lengths treated as compile-time constants (enabling stronger graph rewrites) without restructuring your model, use the `freeze_dims_and_data` model transform, which converts `pm.Data` containers into true constants before compilation. An MCMC sampler calls its compiled function many times, proposing new parameter points and reading back log-probability (and gradient, when needed) 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. From 8bba5b818427acc622355adde343319de2158991 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Fri, 26 Jun 2026 19:09:54 +0100 Subject: [PATCH 6/7] docs: broaden intro motivation bullet list Enumerate more of what the guide explains beyond compilation/caching: pm.Data mutability, numerical stability, backend choice, and why vectorized code helps. Co-authored-by: Cursor --- docs/source/learn/core_notebooks/pytensor_compilation.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/source/learn/core_notebooks/pytensor_compilation.md b/docs/source/learn/core_notebooks/pytensor_compilation.md index 6409505c7a..d35f9b6e4a 100644 --- a/docs/source/learn/core_notebooks/pytensor_compilation.md +++ b/docs/source/learn/core_notebooks/pytensor_compilation.md @@ -3,10 +3,14 @@ **Author:** [Benjamin Vincent](https://github.com/drbenvincent) -In your day-to-day modeling work, it is possible to write models exclusively in PyMC and never know about, or write any, PyTensor code. But PyTensor is a _crucial_ part of the modeling stack, and does a lot of work on your behalf every time you sample. Understanding what it is and how it works can help explain two things PyMC users often wonder about: +In your day-to-day modeling work, it is possible to write models exclusively in PyMC and never know about, or write any, PyTensor code. But PyTensor is a _crucial_ part of the modeling stack, and does a lot of work on your behalf every time you sample. Understanding what it does helps with questions like: - **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. +- **Why should I use `pm.Data` instead of passing arrays directly?** Because `pm.Data` lets you update your data without recompiling — critical for iterative workflows. +- **Why does PyMC rarely produce `NaN` or `inf` from ordinary log-probability expressions?** Numerical rewrites fix fragile formulas before anything runs. +- **Which backend should I choose, and what does that even mean?** This guide explains the options and when each matters. +- **Why does writing vectorized code matter even when the backend isn't Python?** Because vectorized expressions produce simpler graphs that the rewriter can optimise and translate more effectively. This guide gives you a mental model of the pipeline and the choices available, without assuming you've ever written a line of PyTensor. From 9be7dc021172e121a4b510d50259c2dcc8b32e90 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Vincent" Date: Fri, 3 Jul 2026 20:25:13 +0100 Subject: [PATCH 7/7] Update docs/source/learn/core_notebooks/pytensor_compilation.md Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com> --- docs/source/learn/core_notebooks/pytensor_compilation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/learn/core_notebooks/pytensor_compilation.md b/docs/source/learn/core_notebooks/pytensor_compilation.md index d35f9b6e4a..9d83bcfd67 100644 --- a/docs/source/learn/core_notebooks/pytensor_compilation.md +++ b/docs/source/learn/core_notebooks/pytensor_compilation.md @@ -3,7 +3,7 @@ **Author:** [Benjamin Vincent](https://github.com/drbenvincent) -In your day-to-day modeling work, it is possible to write models exclusively in PyMC and never know about, or write any, PyTensor code. But PyTensor is a _crucial_ part of the modeling stack, and does a lot of work on your behalf every time you sample. Understanding what it does helps with questions like: +In your day-to-day modeling work, it is possible to write models exclusively in PyMC and never know about, let alone write, PyTensor code. Despite remaining in the background, the work Pytensor does on your behalf every time you sample makes it a crucial part of modeling stack. Understanding how it works helps with questions like: - **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.