Add PyMC v5 model implementations for 108/120 posteriors by twiecki · Pull Request #318 · stan-dev/posteriordb

twiecki · 2026-03-13T11:00:15Z

Summary

This PR adds PyMC v5 implementations for 104 out of 120 Stan models in posteriordb, along with comprehensive test infrastructure to verify correctness.

Every transpiled model has been validated for gradient equivalence against the original Stan implementation via BridgeStan. Gradients are invariant to normalization constants, making them a stricter correctness check than log-density comparisons alone. Additionally, log-density values are verified to match up to a constant offset (accounting for Stan's convention of dropping normalizing constants).

All models follow idiomatic PyMC v5 patterns:

Vectorized operations (no Python for-loops over data dimensions)
pytensor.scan for true sequential dependencies (GARCH, state-space, HMMs)
Proper use of pm.HalfNormal, pm.HalfCauchy etc. for lower-bounded parameters
0-based indexing with explicit conversion from Stan's 1-based convention
Consistent make_model(data: dict) -> pm.Model interface

Model coverage

Category	Models	Examples
Linear/GLM	30+	`blr`, `earn_height`, `logistic_regression_rhs` (regularized horseshoe)
Hierarchical	20+	`radon_`, `election88_full`, `eight_schools_`
Time series	10+	`garch11`, `arK`, `prophet`
IRT	4	`2pl_latent_reg_irt`, `hier_2pl`, `grsm_latent_reg_irt`, `irt_2pl`
HMM	3	`hmm_example`, `hmm_drive_0` (forward algorithm)
Capture-recapture	5	`Mh_model`, `Mth_model`, `Mtbh_model`, `Rate_*`
ODE	2	`lotka_volterra`, `one_comp_mm_elim_abs`
Mixture	3	`normal_mixture`, `low_dim_gauss_mix`

Remaining 16 models

These require specialized handling beyond the current transpiler capabilities:

HMMs with complex state spaces (hmm_gaussian, hmm_drive_1, iohmm_reg)
Gaussian processes (hierarchical_gp, kronecker_gp) — need pm.gp API
Topic models (ldaK2, ldaK5) — discrete marginalization
Epidemiological ODEs (sir, covid19imperial_v2/v3)
RBMs (nn_rbm1bJ10, nn_rbm1bJ100)

Test infrastructure

Two test suites are included:

test_transpiled_models.py — Compares log-density values between PyMC and BridgeStan at multiple parameter points. Allows for constant offsets from normalization conventions.
tests/test_pymc_gradients.py — Compares gradients of the log-density, which are invariant to additive constants. This is the primary correctness verification since identical gradients guarantee identical posterior geometry.

Both test suites auto-discover all models that have both a Stan and PyMC implementation, so new models are automatically tested.

Review

I realize this is a lot of code to review at once. I'm happy to:

Split this into individual per-model PRs if that's preferred
Tag PyMC core devs for review of the PyMC idioms and patterns

Let me know what works best.

Review checklist (104 models)

🤖 Generated with Claude Code

…oled Transpiled from Stan using pymc-rust-ai-compiler's Stan→PyMC transpiler. All models validated against BridgeStan reference logp values. https://claude.ai/code/session_012idBhKFGF4Ju757RqTpMcD

Add PyMC v5 transpiled models (blr, earn_height, wells_dist, radon_pooled)

posterior_database/models/pymc/2pl_latent_reg_irt.py

ricardoV94 · 2026-03-13T11:52:12Z

posterior_database/models/pymc/accel_gp.py

+
+        zgp_sigma_1 = pm.Normal("zgp_sigma_1", mu=0, sigma=1, shape=NBgp_sigma_1)
+
+        # Custom potentials to match Stan's exact truncated Student-t implementation


Why can't it use pm.Truncated?

ricardoV94 · 2026-03-13T11:53:17Z

posterior_database/models/pymc/accel_splines.py

+    import numpy as np
+
+    with pm.Model() as model:
+        # Extract data


In general I like the first model writing approach better, do the data wrangling outside of pm.Model() so I can skip to it faster

posterior_database/models/pymc/accel_splines.py

ricardoV94 · 2026-03-13T11:57:52Z

posterior_database/models/pymc/arma11.py

+        errs, _ = pytensor.scan(
+            fn=step,
+            sequences=[y_tensor[1:], y_tensor[:-1]],
+            outputs_info=[err_0],
+            non_sequences=[mu, phi, theta],
+        )
+        err = pt.concatenate([pt.atleast_1d(err_0), errs])
+
+        # Likelihood: err ~ normal(0, sigma) using Potential
+        log_likelihood = pt.sum(pm.logp(pm.Normal.dist(mu=0, sigma=sigma), err))
+        pm.Potential("likelihood", log_likelihood)


Did it try this approach? https://www.pymc.io/projects/examples/en/latest/time_series/Time_Series_Generative_Graph.html

It can definitely do ar and ma, so arma should be fine? Older POC: https://gist.github.com/ricardoV94/a49b2cc1cf0f32a5f6dc31d6856ccb63#file-pymc_timeseries_ma-ipynb

ricardoV94 · 2026-03-13T11:58:24Z

posterior_database/models/pymc/bones_model.py

+        # Parameters - use Flat priors and add manual normal log prob to match Stan exactly
+        theta = pm.Flat("theta", shape=nChild)
+
+        # Add manual prior to match Stan's normal(0, 36) exactly


ricardoV94 · 2026-03-13T11:59:24Z

Too eager to jump to pm.Potential

ricardoV94 · 2026-03-13T12:03:30Z

The discrete marginalization /HMMM would be great test cases for missing functionality in pymc_extras.marginalize. But I would definitely split that work

twiecki · 2026-03-13T12:04:14Z

The discrete marginalization /HMMM would be great test cases for missing functionality in pymc_extras.marginalize. But I would definitely split that work

Yes, that's actually running right now.

twiecki · 2026-03-13T12:06:12Z

The discrete marginalization /HMMM would be great test cases for missing functionality in pymc_extras.marginalize. But I would definitely split that work

oh, you mean currently marginalize can't solve those?

ricardoV94 · 2026-03-13T12:17:32Z

The discrete marginalization /HMMM would be great test cases for missing functionality in pymc_extras.marginalize. But I would definitely split that work

oh, you mean currently marginalize can't solve those?

Dunno. It's restricted to what graphs it allow marginalization. I didn't take a look, I assumed you excluded all these cases, reading the top message.

Adds run_compile_to_rust.py batch script that uses the transpailer agentic loop (Claude + cargo build + logp validation) to compile PyMC models to optimized Rust logp+gradient implementations. Successfully compiled models: blr, diamonds, kidscore_interaction, kidscore_interaction_c, kidscore_interaction_c2, kidscore_interaction_z. Each model includes generated.rs and optimization trace (results.tsv). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Add PyMC-to-Rust compilation pipeline (6/104 models)

MansMeg · 2026-03-13T18:20:11Z

This is very good news. We have been working to find a way to start to add the models. I think we so far has been adding PyMC models and testing them. @JTorgander, what do you think about this PR? Is there a way to accept them in bulk using our testing process?

Hand-ported Stan models using Stan ILR+softmax simplex transform (Helmert sub-matrix basis) with correct Jacobians. All 4 models pass gradient validation against BridgeStan at rtol=1e-5, atol=1e-6. Models added: - ldaK2: Latent Dirichlet Allocation (K=2 topics) - ldaK5: Latent Dirichlet Allocation (K=5 topics) - hmm_drive_1: Hidden Markov Model with bivariate emissions - hierarchical_gp: Hierarchical Gaussian Process with variance decomposition Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Address reviewer feedback on PR stan-dev#318: - Replace pt.dot()/pm.math.dot() with @ operator (16 files) - Remove constant correction Potentials that don't affect sampling (30 files) - Remove unnecessary shape=1 special cases in accel_splines - Replace Flat+Potential prior pattern with pm.Normal in bones_model These changes make the transpiled models more idiomatic PyMC while preserving gradient equivalence with the original Stan models. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Hand-ported Stan models using Stan ILR+softmax simplex transform (Helmert sub-matrix basis) with correct Jacobians. All 4 models pass gradient validation against BridgeStan at rtol=1e-5, atol=1e-6. Models added: - ldaK2: Latent Dirichlet Allocation (K=2 topics) - ldaK5: Latent Dirichlet Allocation (K=5 topics) - hmm_drive_1: Hidden Markov Model with bivariate emissions - hierarchical_gp: Hierarchical Gaussian Process with variance decomposition Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Porting guide for remaining 16 Stan-to-PyMC models

ricardoV94 · 2026-03-14T10:12:14Z

@twiecki can you ask it to select the models that are idiomatic (no weird helper functions, no potentials) and split those into a separate PR? Also do ask it do all the numpy data wrangling before entering pm.Model

Add 4 PyMC models with simplex parameters

twiecki · 2026-03-15T11:15:22Z

@twiecki can you ask it to select the models that are idiomatic (no weird helper functions, no potentials) and split those into a separate PR? Also do ask it do all the numpy data wrangling before entering pm.Model

#319

Add initvals to 4 models with initialization issues (Flat/HalfFlat priors, ordered transforms, high-dimensional latent params). Priors unchanged so logp tests remain valid. Transpile hmm_drive_1 (forward algorithm HMM). - accel_splines: initval on Flat spline coefficients + Truncated sds - low_dim_gauss_mix: initval=[-1, 1] for ordered mu - lsat_model: initval=zeros for 1000 latent thetas - kidscore_mom_work: initval for Flat beta + HalfFlat sigma - hmm_drive_1: new transpilation with forward algorithm as Potential Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Fix sampling failures in 5 transpiled PyMC models

ricardoV94

some code smell in some of those initval. Ordered usually needs it, but others are already zero by default. I assume your bot is not using model.initial_point but trying to read model.rvs_to_initial_point (or whatever is called directly)

posterior_database/models/pymc/accel_splines.py

posterior_database/models/pymc/kidscore_mom_work.py

Tested each model without initvals to find the minimum set: - accel_splines: only Truncated sds needs initval (Flat defaults to 0) - low_dim_gauss_mix: only ordered mu needs initval - hmm_drive_1: only ordered phi/lambda need initval - kidscore_mom_work: no initvals needed (samples fine with defaults) - lsat_model: no initvals needed (samples fine with defaults) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Remove redundant initvals, keep only where needed

twiecki · 2026-03-17T01:23:12Z

Should be set.

Benchmark PyMC (nutpie/numba) vs Stan (cmdstan) on all posteriordb models. Results: PyMC faster on 52/101 models, geometric mean 1.30x speedup. Includes benchmark script, per-model results, and visualization plots. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Switch primary metric from raw sampling time to total wall-clock time (compile + sample) per effective sample. PyMC wins 85/101 models (1.90x geo mean) on this end-to-end efficiency metric. Add separate plots for sec/ESS sampling-only, sec/ESS total, raw time, and total time. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Replace pm.Flat + pm.Potential anti-pattern with proper distributions in surgical_model, logmesquite_logva, logmesquite_logvas, and arma11. surgical_model now converges (was Rhat=4.03 with 3615 divergences). Updated results: PyMC wins 87% of models on total sec/ESS (was 84%), geometric mean advantage 2.04x (was 1.90x). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

twiecki · 2026-03-19T05:15:07Z

Closing in favor of #319 and #320.

claude and others added 30 commits March 11, 2026 00:23

Add PyMC v5 transpiled models: blr, earn_height, wells_dist, radon_po…

e867a47

…oled Transpiled from Stan using pymc-rust-ai-compiler's Stan→PyMC transpiler. All models validated against BridgeStan reference logp values. https://claude.ai/code/session_012idBhKFGF4Ju757RqTpMcD

Merge pull request #1 from twiecki/pymc-models

daa3110

Add PyMC v5 transpiled models (blr, earn_height, wells_dist, radon_pooled)

Merge branch 'stan-dev:master' into master

944b7fd

Add PyMC port: Rate_1_model

c9d0f6e

Add PyMC port: nes_logit_model

3bee9dc

Add PyMC port: kidscore_momhs

874d835

Add PyMC port: kidscore_momiq

7518764

Add PyMC port: logearn_height

222681c

Add PyMC port: wells_dist100_model

f445daf

Add PyMC port: Rate_3_model

25c1ba5

Add PyMC port: kidscore_momhsiq

d369dac

Add PyMC port: wells_dist100ars_model

fe0443c

Add PyMC port: low_dim_gauss_mix_collapse

17542db

Add PyMC port: normal_mixture_k

4fa9f97

Add PyMC port: low_dim_gauss_mix

5ca66a1

Add PyMC port: wells_dae_model

8152dc4

Add PyMC port: wells_interaction_model

8a8ce7e

Add PyMC port: logmesquite_logvolume

b2b3ee0

Add PyMC port: kidscore_interaction

fc4244d

Add PyMC port: mesquite

74609d5

Add PyMC port: kidscore_mom_work

e23be1e

Add PyMC port: wells_interaction_c_model

0345ead

Add PyMC port: kidscore_interaction_c

75f03d8

Add PyMC port: kidscore_interaction_c2

37e22d7

Add PyMC port: wells_dae_c_model

ab34e52

Add PyMC port: normal_mixture

27b2a89

Add PyMC port: kidscore_interaction_z

3a4441c

Add PyMC port: wells_daae_c_model

95c1cee

Add PyMC port: Rate_5_model

b3e62f3

Add PyMC port: logmesquite_logva

571b820

ricardoV94 reviewed Mar 13, 2026

View reviewed changes

twiecki and others added 2 commits March 14, 2026 01:28

Merge pull request #10 from twiecki/rust-compilation-batch1

fcd9bb1

Add PyMC-to-Rust compilation pipeline (6/104 models)

twiecki and others added 4 commits March 14, 2026 14:59

Merge pull request #9 from twiecki/porting-guide-remaining-16

ec56d05

Porting guide for remaining 16 Stan-to-PyMC models

Merge pull request #11 from twiecki/pymc-simplex-models

af19729

Add 4 PyMC models with simplex parameters

twiecki changed the title ~~Add PyMC v5 model implementations for 104/120 posteriors~~ Add PyMC v5 model implementations for 108/120 posteriors Mar 15, 2026

twiecki and others added 3 commits March 17, 2026 11:16

Fix import: pymc_rust_compiler → transpailer

33ce5aa

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Merge pull request #12 from twiecki/fix/sampling-failures-initvals

eb3d357

Fix sampling failures in 5 transpiled PyMC models

ricardoV94 reviewed Mar 16, 2026

View reviewed changes

posterior_database/models/pymc/accel_splines.py Outdated Show resolved Hide resolved

posterior_database/models/pymc/accel_splines.py Outdated Show resolved Hide resolved

posterior_database/models/pymc/kidscore_mom_work.py Outdated Show resolved Hide resolved

twiecki and others added 2 commits March 17, 2026 12:52

Merge pull request #13 from twiecki/fix/sampling-failures-initvals

9281b68

Remove redundant initvals, keep only where needed

twiecki and others added 4 commits March 17, 2026 17:27

Add benchmark analysis: PyMC/nutpie vs Stan on 101 posteriordb models

5f7fb6c

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

twiecki closed this Mar 19, 2026


		zgp_sigma_1 = pm.Normal("zgp_sigma_1", mu=0, sigma=1, shape=NBgp_sigma_1)

		# Custom potentials to match Stan's exact truncated Student-t implementation

Conversation

twiecki commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Model coverage

Remaining 16 models

Test infrastructure

Review

Review checklist (104 models)

Uh oh!

Uh oh!

ricardoV94 Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Mar 13, 2026

Uh oh!

ricardoV94 commented Mar 13, 2026

Uh oh!

twiecki commented Mar 13, 2026

Uh oh!

twiecki commented Mar 13, 2026

Uh oh!

ricardoV94 commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MansMeg commented Mar 13, 2026

Uh oh!

ricardoV94 commented Mar 14, 2026

Uh oh!

twiecki commented Mar 15, 2026

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

twiecki commented Mar 17, 2026

Uh oh!

twiecki commented Mar 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

twiecki commented Mar 13, 2026 •

edited

Loading

ricardoV94 Mar 13, 2026 •

edited

Loading

ricardoV94 commented Mar 13, 2026 •

edited

Loading