Statespace test cleanup

Author: jessegrabowski

tests/statespace/core/test_statespace_JAX.py

@@ -24,7 +24,7 @@
 from tests.statespace.test_utilities import load_nile_test_data
 
 pytest.importorskip("jax")
-pytest.importorskip("numpyro")
+pytest.importorskip("nutpie")
 
 
 floatX = pytensor.config.floatX
@@ -78,7 +78,8 @@ def idata(pymc_mod, rng, mock_pymc_sample):
                 tune=1,
                 chains=1,
                 random_seed=rng,
-                nuts_sampler="numpyro",
+                nuts_sampler="nutpie",
+                nuts_sampler_kwargs={"backend": "jax", "gradient_backend": "jax"},
                 progressbar=False,
             )
         with freeze_dims_and_data(pymc_mod):
@@ -101,7 +102,8 @@ def idata_exog(exog_pymc_mod, rng, mock_pymc_sample):
                 tune=1,
                 chains=1,
                 random_seed=rng,
-                nuts_sampler="numpyro",
+                nuts_sampler="nutpie",
+                nuts_sampler_kwargs={"backend": "jax", "gradient_backend": "jax"},
                 progressbar=False,
             )
         with freeze_dims_and_data(pymc_mod):
@@ -123,8 +125,7 @@ def test_no_nans_in_sampling_output(ss_mod, group, matrix, idata):
 @pytest.mark.parametrize("kind", ["conditional", "unconditional"])
 def test_sampling_methods(group, kind, ss_mod, idata, rng):
     f = getattr(ss_mod, f"sample_{kind}_{group}")
-    with pytest.warns(UserWarning, match="The RandomType SharedVariables"):
-        test_idata = f(idata, random_seed=rng)
+    test_idata = f(idata, random_seed=rng)
 
     if kind == "conditional":
         for output in ["filtered", "predicted", "smoothed"]:
@@ -142,10 +143,9 @@ def test_sampling_methods(group, kind, ss_mod, idata, rng):
 def test_forecast(filter_output, ss_mod, idata, rng):
     time_idx = idata.posterior.coords["time"].values
 
-    with pytest.warns(UserWarning, match="The RandomType SharedVariables"):
-        forecast_idata = ss_mod.forecast(
-            idata, start=time_idx[-1], periods=10, filter_output=filter_output, random_seed=rng
-        )
+    forecast_idata = ss_mod.forecast(
+        idata, start=time_idx[-1], periods=10, filter_output=filter_output, random_seed=rng
+    )
 
     assert forecast_idata.coords["time"].values.shape == (10,)
     assert forecast_idata.forecast_latent.dims == ("chain", "draw", "time", "state")

tests/statespace/filters/test_distributions.py

@@ -93,7 +93,7 @@ def ss_mod_no_me():
     return ss_mod
 
 
-@pytest.mark.parametrize("kfilter", filter_names, ids=filter_names)
+@pytest.mark.parametrize("kfilter", filter_names)
 def test_loglike_vectors_agree(kfilter, pymc_model):
     # TODO: This test might be flakey, I've gotten random failures
     ss_mod = structural.LevelTrendComponent(order=2).build(

tests/statespace/models/test_SARIMAX.py

@@ -432,15 +432,7 @@ def test_SARIMA_with_exogenous(rng, mock_sample):
     obs_intercept = ss_mod.ssm["obs_intercept"]
     assert obs_intercept.type.shape == (None, ss_mod.k_endog)
 
-    intercept_fn = pytensor.function(
-        inputs=list(explicit_graph_inputs(obs_intercept)), outputs=obs_intercept
-    )
     data_val = rng.normal(size=(100, 2)).astype(floatX)
-    beta_val = rng.normal(size=(2,)).astype(floatX)
-
-    intercept_val = intercept_fn(data_val, beta_val)
-    np.testing.assert_allclose(intercept_val, intercept_fn(data_val, beta_val))
-
     data_df = pd.DataFrame(
         rng.normal(size=(100, 1)),
         index=pd.date_range(start="2020-01-01", periods=100, freq="D"),