Pymc Inference Template

PyMC inference template

The reference layout every Bayesian bundle in ManagerPack copies. The goal is consistency: same project structure, same MLflow conventions, same posterior reload path, same notebook style. Whatever the actual model, the plumbing is identical.

The worked example is "is this coin fair?" — a Beta-Binomial conjugate model fit with NUTS. The model is trivial, the plumbing is the point.

This is the Bayesian counterpart to template-sklearn-pipeline. Both templates consume studio/data/coin-flip.parquet. The contrast is the selling point: sklearn gives you a point estimate; PyMC gives you a posterior distribution and lets you answer decision-relevant questions like "what's the probability the coin is fair?"

Project layout

<bundle>/
├── README.md             # what this bundle does + how to run it
├── SKILL.md              # this file (or specialized for the bundle)
├── src/
│   ├── train.py          # PyMC fit + ArviZ diagnostics + MLflow logging
│   ├── predict.py        # load idata from MLflow and answer questions
│   └── plots.py          # ArviZ-based plot helpers
├── notebooks/
│   └── <name>_demo.py    # marimo notebook with conjugate / interactive widgets
└── mlruns/               # MLflow tracking store (gitignored)

The data lives outside the bundle, in studio/data/<problem>.parquet, generated by datagen <problem>. Bundles never carry their own data — they consume parquet from the studio's shared data directory.

Conventions

Model definition

Use a with pm.Model() context manager. Always:

• Give every random variable an explicit name (p, not auto-generated)
• Document the prior choice with a comment
• Use sufficient statistics when they exist (e.g. pm.Binomial("flips", n=N, p=p, observed=k) rather than N independent Bernoulli observations). The likelihood is identical but the sampler is much faster.
• Note in a comment whether the model is conjugate; if so, mention the closed-form solution as a sanity check

with pm.Model() as model:
    p = pm.Beta("p", alpha=prior_alpha, beta=prior_beta)
    pm.Binomial("flips", n=n_flips, p=p, observed=n_heads)

Sampling

Always:

• Use NUTS by default (pm.sample() picks NUTS automatically)
• Set random_seed= for reproducibility
• Use 4 chains so R-hat is meaningful
• Disable progress bar in scripted runs (progressbar=False)
• Tune ≥ 1000, draws ≥ 2000 unless you have a reason to deviate

with model:
    idata = pm.sample(
        draws=2000,
        tune=1000,
        chains=4,
        random_seed=seed,
        progressbar=False,
    )

MLflow logging

Every run logs:

The recovery metrics matter most:

• posterior_mean — point summary
• p_recovery_error = |posterior_mean - true_value|
• true_p_in_94_hdi — boolean, did the credible interval contain the truth?

If the truth is always in the HDI across many runs at the right rate (≈94% of the time), the model is well-calibrated.

Posterior persistence

Save the InferenceData as a NetCDF artifact:

idata_path = tmp_dir / "idata.nc"
idata.to_netcdf(idata_path)
mlflow.log_artifact(str(idata_path), artifact_path="posterior")

There is no mlflow.pymc.log_model like there is for sklearn. The NetCDF artifact + the script that produced it are the canonical serialization. NetCDF is the standard format for ArviZ inference data and round-trips cleanly.

Loading the posterior elsewhere

import arviz as az
import mlflow

mlflow.set_tracking_uri(f"file:{template_dir / 'mlruns'}")
client = mlflow.tracking.MlflowClient()
local_path = client.download_artifacts(run_id, "posterior/idata.nc")
idata = az.from_netcdf(local_path)

The full posterior comes back as one InferenceData object. No need for the original training code.

Diagnostics — never skip these

Always check before trusting a posterior:

• R-hat ≤ 1.01 for every parameter (chains have mixed)
• ESS bulk ≥ 400 per chain (effective sample size is usable)
• No divergences (or document why if there are some)
• Trace plot looks like fuzzy caterpillars (visual sanity check)

Log these as metrics so they're queryable in the MLflow UI. A bundle that ships a model with r_hat > 1.01 is broken.

Plots

Use src/plots.py to generate ArviZ figures. Always include:

• az.plot_posterior(idata, var_names=[...]) with the ground-truth marker overlaid (when known)
• az.plot_trace(idata, var_names=[...]) for chain mixing diagnosis
• A custom prior-vs-posterior overlay using scipy.stats for any conjugate models — this makes the data's effect on the prior obvious

Marimo notebook (always marimo, never Jupyter)

Every bundle ships a notebooks/<name>_demo.py marimo notebook that:

1. Loads the latest run from mlruns/
2. Loads the posterior via arviz.from_netcdf
3. Shows the posterior summary (mean, sd, HDI, R-hat, ESS)
4. Interactive prior exploration via mo.ui.slider. For conjugate models, use the closed-form posterior for instant updates instead of re-sampling. For non-conjugate models, fall back to pre-computing a grid of priors or accepting the slower MCMC re-fit.
5. Shows decision-relevant probabilities (e.g. P(fair))
6. Compares recovered values to ground truth from the sidecar

Validate with marimo check notebooks/<name>_demo.py. Test rendering headlessly with marimo export html notebooks/<name>_demo.py -o /tmp/x.html to confirm all cells execute without errors before shipping.

Running the worked example

# 1. Generate the dataset
datagen coin-flip --n 500 --p 0.7 --seed 42

# 2. Fit and log to MLflow
cd studio/templates/pymc-inference
python src/train.py

# 3. Verify recovery
# stdout shows posterior mean ≈ 0.71, |recovery error| ≈ 0.01
# and "true p in 94% HDI: True"

# 4. Load the posterior and answer questions
python src/predict.py --run-id <id>

# 5. Open the demo notebook and explore the conjugate prior
marimo edit notebooks/coin_flip_demo.py

The "is it fair?" answer

datagen coin-flip --p 0.5 --seed 99 --output /tmp/fair.parquet
python src/train.py --data /tmp/fair.parquet
# → P(p∈[0.5±0.05]) ≈ 0.6  ("60% chance the coin is fair within ±5%")

datagen coin-flip --p 0.7 --seed 42 --output /tmp/biased.parquet
python src/train.py --data /tmp/biased.parquet
# → P(p∈[0.5±0.05]) ≈ 0.0  ("definitely not fair")

This is the kind of decision-relevant output a sklearn point estimate cannot give you. Bundles built on this template should always expose at least one decision-relevant probability metric in the same spot.

Specializing this template for a new bundle

1. Copy studio/templates/pymc-inference/ to studio/scratch/<bundle-name>/ (or directly to bundles/<bundle-name>/ once the design is settled)
2. Replace the dataset reference: change coin-flip to your problem name
3. Update build_model() to use the right priors and likelihood
4. Update the diagnostics and decision-relevant metrics
5. Update the plots in src/plots.py for the new domain
6. Update the marimo notebook with the right interactive widgets — for non-conjugate models, you may need to pre-compute a grid of priors instead of relying on closed-form updates
7. Update SKILL.md with the bundle-specific story