Fix stale trace count in test_mcmc_parallel_chain and run it in CI (#2223)

* Fix stale trace count in test_mcmc_parallel_chain and run it in CI

PR #1910 merged postprocess_fn into the fori_collect loop, adding two
extra model traces for deterministic values, and updated
test_mcmc_one_chain accordingly (4 -> 5) but left test_mcmc_parallel_chain
asserting 4. The parallel-chain path has the same two extra calls, so the
test fails with `assert 5 == 4` on any host with two or more devices.

This went unnoticed because CI only runs test/test_compile.py without
XLA_FLAGS=--xla_force_host_platform_device_count=2, so the skipif marker
always skips the parallel-chain test. Update the assertion to 5 and add
the file to the multi-device "Test chains" batch so it actually runs.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* pyproject.toml: add enough `dev` dependencies to get `uv run make test` passing

Signed-off-by: Eli Sennesh <elisennesh@astera.org>

---------

Signed-off-by: Eli Sennesh <elisennesh@astera.org>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: esennesh

.github/workflows/ci.yml

@@ -148,6 +148,7 @@ jobs:
           XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/contrib/test_tfp.py -k "chain"
           XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/infer/test_hmc_gibbs.py -k "chain"
           XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/infer/test_mcmc.py -k "chain or pmap or vmap"
+          XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/test_compile.py -k "chain"
       - name: Test custom prng
         run: |
           JAX_ENABLE_CUSTOM_PRNG=1 uv run pytest -vs test/infer/test_mcmc.py

pyproject.toml

@@ -49,11 +49,19 @@ dev = [
     "jaxns>=2.6.3,<=2.6.9",
     "matplotlib",
     "optax>=0.0.6",
+    "pandas>=3.0.3",
     "pylab-sdk",                         # jaxns dependency
+    "pyro-api>=0.1.2",
+    "pytest>=9.1.1",
     "pytest-cov",
     "pyyaml",                            # flax dependency
     "requests",                          # pylab dependency
+    "ruff>=0.15.22",
+    "scikit-learn>=1.9.0",
+    "seaborn>=0.13.2",
     "tfp-nightly",
+    "ty>=0.0.57",
+    "wordcloud>=1.9.6",
 ]
 ci = [
     "coverage>=7.13.5",

test/test_compile.py

@@ -53,7 +53,11 @@ def test_mcmc_parallel_chain(deterministic):
     mcmc.get_samples()
 
     if deterministic:
-        assert GLOBAL["count"] == 4
+        # As in test_mcmc_one_chain, we have two extra calls to the model
+        # to get deterministic values:
+        # 1. transform the init state
+        # 2. transform state during the loop
+        assert GLOBAL["count"] == 5
     else:
         assert GLOBAL["count"] == 3