Numpyro Python

NumPyro Python

Build a pure-shape-stable probabilistic program, make independence explicit, thread JAX randomness correctly, choose inference from model structure, and verify both sampler diagnostics and predictive behavior.

Boundary

This is the skill for the official pyro-ppl/numpyro project. “numpypro” in a request is interpreted as NumPyro only when Bayesian/JAX context confirms the intent. Use PyMC for a PyTensor model, Pyro guidance for PyTorch, and ArviZ when the task starts from completed draws.

Know the objects

Site names are the interface joining model, guide, posterior sample dictionaries, and predictive execution. Shapes are part of that interface. Read the program and shape contract before adding a plate, .to_event, vectorization, or JIT.

Ordered workflow

1. State the estimand, likelihood, observation unit, independence structure, missing-data policy, latent support, and prediction target.
2. Write the model with y=None or equivalent for predictive execution. Give every site a stable unique name; keep data-dependent Python control flow out of traced/JIT paths.
3. Work out each distribution's batch and event shape on paper. Add plates for actual conditional independence and .to_event only for dependent event axes.
4. Validate input dtype and shape before entering JAX. Enable required numeric precision before creating arrays or compiling.
5. Obtain one root PRNG key, split it for independent operations, and never reuse a consumed key for another stochastic result.
6. Run a seeded trace/shape check, prior predictive simulation, and small simulated-data recovery before expensive inference.
7. Choose MCMC or SVI from posterior geometry, scale, discrete structure, and accuracy requirement. Record all inference and initialization choices.
8. Diagnose MCMC or validate SVI against a trusted small problem. Generate posterior predictive samples with a fresh key and obs=None.
9. Convert to labeled InferenceData when chain/draw diagnostics, persistence, or model comparison are required; supply coords/dims rather than guessing.

Decision map

Read inference and JAX execution before tuning chains, guides, devices, precision, JIT, or vectorization.

Canonical model and MCMC anchor

import jax.numpy as jnp
from jax import random
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive


def model(x: jnp.ndarray, y: jnp.ndarray | None = None) -> None:
    beta = numpyro.sample("beta", dist.Normal(0, 1).expand((x.shape[1],)).to_event(1))
    intercept = numpyro.sample("intercept", dist.Normal(0, 1.5))
    with numpyro.plate("observation", x.shape[0]):
        numpyro.sample("outcome", dist.Bernoulli(logits=intercept + x @ beta), obs=y)


def fit(x: jnp.ndarray, y: jnp.ndarray, seed: int):
    run_key, predictive_key = random.split(random.key(seed))
    mcmc = MCMC(NUTS(model), num_warmup=1_000, num_samples=1_000, num_chains=4)
    mcmc.run(run_key, x, y)
    samples = mcmc.get_samples(group_by_chain=True)
    predictions = Predictive(model, posterior_samples=mcmc.get_samples())(predictive_key, x, None)
    return mcmc, samples, predictions

group_by_chain=True preserves diagnostic structure. Draw counts are examples, not universal defaults; calibrate them to ESS/MCSE and runtime. Inspect the installed JAX key constructor and NumPyro signatures before copying this anchor.

Non-negotiable rules

• Never reuse a PRNG key for independent random operations. Split once at the ownership boundary and pass subkeys explicitly.
• Never solve a plate error by deleting the plate. Reconcile data axes, distribution batch/event shapes, and model independence.
• Site names must be unique across an execution and stable between model, guide, samples, and predictions. A guide must cover required latent sites exactly as intended.
• JIT traces shapes/dtypes and static Python structure. Avoid data-dependent Python branches, mutation, hidden global randomness, and changing site sets.
• More SVI steps do not prove a good variational family. Inspect loss stability, multiple initializations, predictive fit, and small-problem comparison.
• More MCMC samples do not repair divergences, funnels, bad initialization, or non-identifiability. Reparameterize and standardize before merely tuning.
• Device parallelism and chain method depend on installed JAX devices and versions. Inspect, do not assume GPU/TPU or four parallel devices.
• Prediction must omit observed values and use a fresh key. Check returned site shapes and preserve chain/draw identity when exporting diagnostics.

Inspect the installed version, then use testing and version grounding for static trace checks, shape assertions, reproducibility tests, and installed evidence.

Completion gate

Do not declare completion until every site has unique stable naming; plate, batch, and event shapes are asserted; PRNG ownership has no reuse; prior predictive and simulated recovery checks pass; inference choice is justified; MCMC diagnostics or SVI approximation checks have no unresolved critical failure; predictive execution uses obs=None and a fresh key; JIT inputs have a documented shape/precision policy; and unavailable devices, package versions, or skipped statistical checks are reported.

References

• Program and shape contract
• Inference and JAX execution
• Testing and version grounding