Run build_smc_logfns and SMC prior particles consistently in unconstrained space

[junpenglao]

Summary

The documented contract is that SMC positions are unconstrained: postprocess_fn maps unconstrained → constrained, and init_position is param_info.z. Three pieces broke that contract:

• logprior_fn evaluated position as constrained values through log_density, with no transform and no Jacobian.
• loglikelihood_fn = −potential_fn − logprior_fn therefore mixed the two conventions, since numpyro's potential_fn takes unconstrained values.
• build_prior_sample_fn's Predictive fallback returned constrained particles.

The fix:

• logprior_fn(z) = −blocked_potential_fn(z), where blocked_potential_fn is the potential of the prior-only (observed sites blocked) model, built by its own initialize_model. It uses the same unconstrained convention and Jacobian as the joint potential_fn, so the Jacobians cancel and loglikelihood_fn(z) is exactly log p(data | constrain(z)).
• The prior particles are mapped through each site's inverse bijector, so they come out unconstrained.
• The docstrings are corrected.
• The joint initialize_model still receives rng_key unchanged. init_position is bit-identical to main, checked on 2 models × 3 seeds.

Impact on results: none for existing end-to-end SMC runs. logprior_fn + loglikelihood_fn ≡ −potential_fn held by construction, so the λ=1 target was always correct. The old constrained particles were also self-consistent with the old logprior_fn at λ=0. As a result, adaptive tempered SMC with an RWM kernel recovered the NUTS posterior mean (within about 0.005) on a toy model with both old and new code. The bug affected consumers that use logprior_fn, loglikelihood_fn or the particles on their own, and kernels that need genuinely unconstrained positions. Of the committed SMC recipes, gmm_25 and neals_funnel use the analytic-sampler path and logistic_synthetic has only real-valued latents, so none of them hit the broken branch.

Found during the adversarial review of #265 (trax local Issue#1266).

Test plan

• New tests/numpyro/test_smc_logfns_space.py, a toy model with sigma ~ HalfNormal(2) and y ~ Normal(0, sigma).
  • Tests that fail on main:
    • logprior_fn and loglikelihood_fn against an independent analytic reference at 7 unconstrained z. On main, loglikelihood_fn is off at all 7, e.g. z=2: −18.92 vs −14.60.
    • Prior particles include negative values, and their postprocessed mean recovers E[HalfNormal(2)] = 1.60. On main it is 14.11.
  • Structural invariants, which pass on both: logprior + loglik == −potential, and the blocked model's latent sites match the joint model's.
  • Of the 6 new tests, 4 fail on main.
• Existing SMC tests: test_emit_smc_script.py, test_smc_generated_lifecycle.py, test_generated_smc.py, tests/smc/, test_api_pins_smc.py, test_registry.py (162 passed together with the new file).
• pytest tests -m fast -n 2: 2089 passed, 5 skipped.
• pre-commit clean.
• CI green.

🤖 Generated with Claude Code