---

name: modeling-primitives description: "Write and inspect NumPyro model programs with primitives, plates, handlers, JAX PRNG keys, control flow, and model tracing utilities." disable-model-invocation: true metadata: disco-role: operating license: Apache 2.0

modeling-primitives

Use this sub-skill when the task is to write, inspect, or debug a NumPyro model program before choosing an inference algorithm: sample/param sites, plate shapes, effect handlers, JAX PRNG keys, conditioning, masking, reparameterization handlers, model traces, control flow, and model rendering.

Quick workflow

1. Write a pure JAX-compatible model. Use references/modeling-primitives.md for numpyro.sample, param, plate, deterministic, factor, PRNG behavior, and functional-style constraints.
2. Separate distribution shape from model-site shape. First validate the distribution object with ../distributions-transforms/, then place observation axes in plate here.
3. Use handlers for alternate interpretations. Use references/effect-handlers.md for seed, trace, condition, substitute, replay, block, mask, scale, reparam, do, and scope.
4. Inspect before inference. Use references/model-inspection.md to trace a model, check site metadata, render a dependency graph when graphviz is installed, or summarize shapes.
5. Use JAX control flow for dynamic programs. Prefer numpyro.contrib.control_flow.scan and cond when Python loops or branches interact with JIT, vectorization, or time-series latent structure.
6. Run the smoke script. scripts/model_trace_smoke.py checks deterministic seeding, conditioning, and trace metadata without MCMC/SVI.
7. Triage symptoms. Use references/troubleshooting.md for missing RNG keys, duplicate site names, plate collisions, tracer errors, masked observation confusion, and observed/conditioned site mistakes.

Route elsewhere when

• The task is choosing or fixing a distribution constructor, support constraint, transform, or log_prob shape: use ../distributions-transforms/.
• The task is running NUTS/HMC or interpreting divergences, ESS, r_hat, or posterior predictive output from MCMC: use ../mcmc-diagnostics/.
• The task is fitting a guide, choosing an ELBO, using autoguides, or handling SVI losses: use ../svi-autoguides/.
• The task depends on optional contrib modules such as Flax/Equinox wrappers, Funsor, TFP, nested sampling, HSGP, or SteinVI: use ../advanced-contrib/.

Common trigger phrases

• "Trace this model", "condition observations", "seed a model", "why does sample need a key?"
• "Fix plate shape", "use subsampling", "mask observations", "make this model JAX-compatible"
• "Render a model graph", "inspect model relations", "convert a loop to scan"
• "Use reparam handler", "hide a site", "substitute parameter values"

Bundled materials

• Modeling primitives explains model-site APIs and JAX/PRNG rules.
• Effect handlers maps handler names to concrete use-cases.
• Model inspection shows trace, render, relation, and shape workflows.
• Troubleshooting covers predictable model-construction failures.
• Model trace smoke is a small assertion-backed diagnostic script.