DO NOT MERGE: pip-install nutpie from main in all CI jobs

Needed until conda-forge ships a nutpie release that supports arviz>=1.0.
Until then the conda env can't co-install nutpie and pymc's arviz pin,
so we get nutpie via pip-from-git-main in every CI job. Once a compatible
nutpie is on conda-forge, drop this commit and put `nutpie>=<that version>`
back into `environment-test.yml` and `windows-environment-test.yml`.

Author: ricardoV94

.github/workflows/tests.yml

@@ -168,6 +168,8 @@ jobs:
         run: |
           pip install "pytensor @ git+https://github.com/pymc-devs/pytensor.git@v3" --no-deps
           pip install -e . --no-deps
+          # Track nutpie main to catch any incompatibility with pymc early.
+          pip install "nutpie @ git+https://github.com/pymc-devs/nutpie@main"
           python --version
           micromamba list
       - name: Run tests
@@ -218,6 +220,8 @@ jobs:
         run: |
           pip install "pytensor @ git+https://github.com/pymc-devs/pytensor.git@v3" --no-deps
           pip install -e . --no-deps
+          # Track nutpie main to catch any incompatibility with pymc early.
+          pip install "nutpie @ git+https://github.com/pymc-devs/nutpie@main"
           python --version
           micromamba list
       - name: Run tests
@@ -276,6 +280,8 @@ jobs:
         run: |
           pip install "pytensor @ git+https://github.com/pymc-devs/pytensor.git@v3" --no-deps
           pip install -e . --no-deps
+          # Track nutpie main to catch any incompatibility with pymc early.
+          pip install "nutpie @ git+https://github.com/pymc-devs/nutpie@main"
           python --version
           micromamba list
       - name: Run tests
@@ -378,6 +384,8 @@ jobs:
         run: |
           pip install "pytensor @ git+https://github.com/pymc-devs/pytensor.git@v3" --no-deps
           pip install -e . --no-deps
+          # Track nutpie main to catch any incompatibility with pymc early.
+          pip install "nutpie @ git+https://github.com/pymc-devs/nutpie@main"
           python --version
           micromamba list
       - name: Run tests

pymc/sampling/mcmc.py

@@ -908,11 +908,12 @@ def sample(
 
     if nuts_sampler is None:
         # Nutpie must take all the variables and can only compile to Numba or JAX.
-        # When the user asks for a MultiTrace or provides a pymc-backend trace object
-        # (e.g. `ZarrTrace`), route to the pymc sampler.
+        # When the user asks for a MultiTrace, a pymc-only `init`, or provides a
+        # pymc-backend trace object (e.g. `ZarrTrace`), route to the pymc sampler.
         can_use_nutpie = (
             exclusive_nuts
             and NUTPIE_INSTALLED
+            and init == "auto"
             and return_inferencedata
             and trace is None
             and callback is None

tests/sampling/test_mcmc_external.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 
 import unittest.mock as mock
+import warnings
 
 from types import SimpleNamespace
 
@@ -214,3 +215,76 @@ def test_nutpie_end_to_end():
     assert {"posterior", "sample_stats", "observed_data"} <= set(idata.children)
     assert set(idata.posterior.data_vars) == {"mu", "sigma"}
     assert idata.posterior.sizes == {"chain": 2, "draw": 20}
+
+
+class TestExternalSamplerKwargCompat:
+    """Validate how `pm.sample` handles kwargs that external samplers don't fully honor."""
+
+    @pytest.fixture
+    def model(self):
+        with Model() as m:
+            Normal("x")
+        return m
+
+    @pytest.fixture
+    def patched_sampler(self):
+        with (
+            mock.patch("pymc.sampling.mcmc.NUTPIE_INSTALLED", True),
+            mock.patch("pymc.sampling.mcmc._sample_external_nuts") as mock_ext,
+        ):
+            yield mock_ext
+
+    _BASE_KWARGS = {
+        "tune": 2,
+        "draws": 2,
+        "chains": 1,
+        "progressbar": False,
+        "compile_kwargs": {"mode": "NUMBA"},
+    }
+
+    @pytest.mark.parametrize(
+        "extra",
+        [
+            {"return_inferencedata": False},
+            {"trace": object()},
+            {"callback": lambda **kw: None},
+        ],
+        ids=["return_inferencedata_false", "custom_trace", "callback"],
+    )
+    def test_explicit_nutpie_raises_on_incompatible(self, model, patched_sampler, extra):
+        # Filter the separate FutureWarning for return_inferencedata=False; we only
+        # care that pm.sample raises the external-sampler ValueError.
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", ".*return_inferencedata=False.*", FutureWarning)
+            with model, pytest.raises(ValueError, match="`nuts_sampler='nutpie'`"):
+                sample(nuts_sampler="nutpie", **self._BASE_KWARGS, **extra)
+        patched_sampler.assert_not_called()
+
+    def test_explicit_nutpie_warns_on_non_default_init(self, model, patched_sampler):
+        with model, pytest.warns(UserWarning, match="`init='advi'` is ignored"):
+            sample(nuts_sampler="nutpie", init="advi", **self._BASE_KWARGS)
+        patched_sampler.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "extra",
+        [
+            {"return_inferencedata": False},
+            {"trace": object()},
+            {"callback": lambda **kw: None},
+            {"init": "advi"},
+        ],
+        ids=["return_inferencedata_false", "custom_trace", "callback", "non_default_init"],
+    )
+    def test_auto_selection_falls_back_to_pymc(self, model, patched_sampler, extra):
+        # Each of these kwargs disqualifies nutpie auto-selection; pm.sample should
+        # route to the pymc sampler (external stub never called) instead of raising.
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", ".*return_inferencedata=False.*", FutureWarning)
+            with model:
+                try:
+                    sample(**self._BASE_KWARGS, **extra)
+                except Exception:
+                    # The pymc path may error on the stand-in objects (dummy trace, etc.);
+                    # we only care that nutpie wasn't dispatched.
+                    pass
+        patched_sampler.assert_not_called()