NumPyro

This page discharges the per-target obligations of Transpilation correctness for \(\mathsf{T} = \mathrm{NumPyro}\).

Semantics

NumPyro's denotational semantics is the trace semantics of Goodman et al. (2008) / Wingate, Stuhlmüller, Goodman (2011) implemented via the effect-handler architecture of Phan, Pradhan, and Jankowiak (2019). A program is a JAX-traced Python function whose numpyro.sample invocations register sites in a trace. The log-density probe is numpyro.infer.util.log_density; NumPyro returns the joint log-density of the trace.

Unconstrained-space change of variables

NumPyro applies the same change-of-variables principle as Stan but with reparametrization handled by per-site TransformReparam or auto-transformations registered in AutoNormal. For QVR-emitted code at the level the renderer produces, no explicit reparametrization is requested: \(\Psi_{\mathsf{NumPyro}} = \mathrm{id}\) and \(\log|\det J| \equiv 0\). The inference layer may add reparametrizations at runtime, but those are guide-side and do not affect the model-side denotation Theorem 6.1 quantifies over.

Family parameterizations

NumPyro families track PyTorch / Pyro's Distribution hierarchy. The mapping is identity for every QVR family that has a corresponding numpyro.distributions.* class:

  • Normal(μ, σ) ↔ Normal(loc=μ, scale=σ)
  • Dirichlet(α) ↔ Dirichlet(concentration=α) (with jnp.full((K,), α) broadcast when α is scalar)
  • Categorical(p) ↔ Categorical(probs=p)
  • Bernoulli(p) ↔ Bernoulli(probs=p)
  • LogitNormal ↔ LogitNormal(loc=μ, scale=σ)

In each case \(\pi_{F, \mathsf{NumPyro}}\) is the identity and \(c_{F, \mathsf{NumPyro}} = 0\).

Per-construct emit

Sample / observe. numpyro.sample("x", <dist>) for latents; numpyro.sample("y", <dist>, obs=y_data) for observations. Each contributes the documented per-site log-density to the log_density probe.

Plate. Nested with numpyro.plate(name, B): contexts per the documented plate primitive (Phan, Pradhan, Jankowiak 2019 §3.2). The semantics is exactly the conditionally-independent product measure of \(B\) i.i.d. draws (plate discussion).

Marginalize. The NumPyro renderer lowers IRMarginalize to IRSample(latent) + scope body inline (head the marginalization discussion). The latent contributes a normal sample site; the scope's observe site runs with the same trace. Soundness for inference targets that estimate the \(\theta\)-posterior is the projection property of Fritz 2020 Proposition 5.4.

Score / let / return. numpyro.factor("name", expr) for score; deterministic assignment for let; native return for the return clause.

Acceptance

  • Tier 1 pipeline composition. Structural checks cover the emitted model function and the Lower >> NumPyroRenderer >> EmitPretty(python) pipeline.
  • Tier 2 external syntax. Python's AST parser accepts the generated modules in the external-validation matrix.
  • Tier 3 numeric equivalence. For selected fixtures, numpyro.infer.util.log_density is compared with the QVR reference on the shared finite grid and tolerance described in the test contract.