Fix stale trace count in test_mcmc_parallel_chain and run it in CI

[esennesh]

test/test_compile.py::test_mcmc_parallel_chain[True] fails with assert 5 == 4 on any host with jax.device_count() >= 2.

#1910 merged postprocess_fn into the fori_collect loop, which adds two extra model traces when the model has deterministic sites (transform the init state + transform states during the loop). That PR updated test_mcmc_one_chain's expected count from 4 to 5 accordingly, but test_mcmc_parallel_chain was left asserting 4 — and the parallel-chain (pmap) path has the same two extra calls, so the model is now traced 5 times there too.

This went unnoticed for so long because CI never actually runs the test: test/test_compile.py is only executed in the batch without XLA_FLAGS="--xla_force_host_platform_device_count=2", so jax.device_count() is 1 and the skipif marker always skips it.

Bisection confirms this is independent of the JAX version: the test passes at 0e7bd20~1 and fails at 0e7bd20 (#1910) with the same jax 0.4.35, and current master fails identically on jax 0.7.2 through 0.10.2.

Changes made

• Update test_mcmc_parallel_chain's deterministic trace-count assertion from 4 to 5, with the same explanatory comment Merge postprocess_fn into the fori_collect loop #1910 added to test_mcmc_one_chain.
• Add test/test_compile.py -k "chain" to the multi-device "Test chains" CI batch so the test actually runs in CI instead of being silently skipped.

Links to related issues/PRs

• Merge postprocess_fn into the fori_collect loop #1910, which changed the trace count and updated the single-chain test but not the parallel-chain one.

Tests

• XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/test_compile.py -k "chain" — 6 passed (previously test_mcmc_parallel_chain[True] failed with assert 5 == 4).
• pytest test/test_compile.py (single device) — 6 passed, 2 skipped, unchanged.

Dependencies

• None.

🤖 Generated with Claude Code

[fehiepsi]

Thanks Eli!