Install nutpie from main

Author: ricardoV94

.github/workflows/tests.yml

@@ -328,6 +328,9 @@ jobs:
         run: |
           pip install "pytensor @ git+https://github.com/pymc-devs/pytensor.git@v3" --no-deps
           pip install -e . --no-deps
+          # Track nutpie main so we exercise unreleased changes. Expected to be
+          # red until nutpie main drops its `arviz.InferenceData` references.
+          pip install "nutpie @ git+https://github.com/pymc-devs/nutpie@main"
           python --version
           micromamba list
       - name: Run tests

conda-envs/environment-alternative-backends.yml

@@ -11,7 +11,8 @@ dependencies:
 - cloudpickle
 - zarr>=2.5.0,<3
 - numba
-- nutpie >= 0.15.1
+# nutpie is installed from git in the CI workflow (see tests.yml) so we can track
+# the arviz 1.0 -compatible branch until an upstream release ships.
 # Jaxlib version must not be greater than jax version!
 - jax>=0.4.28
 - jaxlib>=0.4.28

tests/sampling/test_mcmc_external.py

@@ -25,9 +25,7 @@
 from pymc.progress_bar import NutpieProgressBarManager
 
 
-# temporarily skip nutpie
-@pytest.mark.parametrize("nuts_sampler", ["pymc", "blackjax", "numpyro"])
-# @pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"])
+@pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"])
 def test_external_nuts_sampler(recwarn, nuts_sampler):
     if nuts_sampler != "pymc":
         pytest.importorskip(nuts_sampler)
@@ -97,9 +95,7 @@ def test_step_args():
     npt.assert_almost_equal(idata.sample_stats.acceptance_rate.mean(), 0.5, decimal=1)
 
 
-# temporarily skip nutpie
-@pytest.mark.parametrize("nuts_sampler", ["pymc", "blackjax", "numpyro"])
-# @pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"])
+@pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"])
 def test_sample_var_names(nuts_sampler):
     if nuts_sampler != "pymc":
         pytest.importorskip(nuts_sampler)