importorskip tests on optional dependencies

Author: jessegrabowski

tests/sampling/test_jax.py

@@ -27,7 +27,6 @@
 import pytensor.tensor as pt
 import pytest
 
-from numpyro.infer import MCMC
 from pytensor.compile import SharedVariable
 from pytensor.graph import graph_inputs
 
@@ -45,6 +44,8 @@
     sample_numpyro_nuts,
 )
 
+MCMC = pytest.importorskip("numpyro.infer.MCMC")
+
 
 def test_jax_PosDefMatrix():
     x = pt.tensor(name="x", shape=(2, 2), dtype="float32")

tests/sampling/test_mcmc_external.py

@@ -76,6 +76,8 @@ def test_external_nuts_sampler(recwarn, nuts_sampler):
 
 
 def test_step_args():
+    pytest.importorskip("numpyro")
+
     with Model() as model:
         a = Normal("a")
         idata = sample(
@@ -91,6 +93,9 @@ def test_step_args():
 
 @pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"])
 def test_sample_var_names(nuts_sampler):
+    if nuts_sampler != "pymc":
+        pytest.importorskip(nuts_sampler)
+
     seed = 1234
     kwargs = {
         "chains": 1,