Add JAX test suite

Author: jessegrabowski

pymc_experimental/tests/statespace/test_kalman_filter.py

@@ -37,11 +37,11 @@
 single_inout = initialize_filter(SingleTimeseriesFilter())
 steadystate_inout = initialize_filter(SteadyStateFilter())
 
-f_standard = pytensor.function(*standard_inout)
-f_cholesky = pytensor.function(*cholesky_inout)
-f_univariate = pytensor.function(*univariate_inout)
-f_single_ts = pytensor.function(*single_inout)
-f_steady = pytensor.function(*steadystate_inout)
+f_standard = pytensor.function(*standard_inout, on_unused_input="ignore")
+f_cholesky = pytensor.function(*cholesky_inout, on_unused_input="ignore")
+f_univariate = pytensor.function(*univariate_inout, on_unused_input="ignore")
+f_single_ts = pytensor.function(*single_inout, on_unused_input="ignore")
+f_steady = pytensor.function(*steadystate_inout, on_unused_input="ignore")
 
 filter_funcs = [f_standard, f_cholesky, f_univariate, f_single_ts, f_steady]
 
@@ -315,3 +315,37 @@ def test_all_covariance_matrices_are_PSD(filter_func, filter_name, n_missing, ob
             atol=ATOL,
             err_msg=f"{name} is not symmetrical",
         )
+
+
+@pytest.mark.parametrize(
+    "filter",
+    [StandardFilter, SingleTimeseriesFilter, CholeskyFilter],
+    ids=["standard", "single_ts", "cholesky"],
+)
+def test_kalman_filter_jax(filter):
+    pytest.importorskip("jax")
+    from pymc.sampling.jax import get_jaxified_graph
+
+    # TODO: Add UnivariateFilter to test; need to figure out the broadcasting issue when 2nd data dim is defined
+
+    p, m, r, n = 1, 5, 1, 10
+    inputs, outputs = initialize_filter(filter(), mode="JAX")
+
+    # Shape of the data must be static for jax to know how long the scan is
+    data = inputs.pop(0)
+    data_specified = pt.specify_shape(data, (n, None))
+    data_specified.name = "data"
+    inputs = [data] + inputs
+
+    outputs = pytensor.graph.clone_replace(outputs, {data: data_specified})
+
+    inputs_np = make_test_inputs(p, m, r, n, rng)
+
+    f_jax = get_jaxified_graph(inputs, outputs)
+    f_pt = pytensor.function(inputs, outputs, mode="FAST_COMPILE")
+
+    jax_outputs = f_jax(*inputs_np)
+    pt_outputs = f_pt(*inputs_np)
+
+    for name, jax_res, pt_res in zip(output_names, jax_outputs, pt_outputs):
+        assert_allclose(jax_res, pt_res, atol=ATOL, rtol=RTOL, err_msg=f"{name} failed!")

pymc_experimental/tests/statespace/test_statespace.py

@@ -94,7 +94,7 @@ def pymc_mod(ss_mod):
         rho = pm.Beta("rho", 1, 1)
         zeta = pm.Deterministic("zeta", 1 - rho)
 
-        ss_mod.build_statespace_graph(data=nile, save_kalman_filter_outputs_in_idata=True)
+        ss_mod.build_statespace_graph(data=nile, save_kalman_filter_outputs_in_idata=False)
         names = ["x0", "P0", "c", "d", "T", "Z", "R", "H", "Q"]
         for name, matrix in zip(names, ss_mod.unpack_statespace()):
             pm.Deterministic(name, matrix)

pymc_experimental/tests/statespace/test_statespace_JAX.py

@@ -0,0 +1,154 @@
+import warnings
+
+import numpy as np
+import pymc as pm
+import pytensor
+import pytensor.tensor as pt
+import pytest
+from pymc.model.transform.optimization import freeze_dims_and_data
+
+from pymc_experimental.statespace.utils.constants import (
+    FILTER_OUTPUT_NAMES,
+    MATRIX_NAMES,
+    SMOOTHER_OUTPUT_NAMES,
+)
+from pymc_experimental.tests.statespace.test_statespace import (  # pylint: disable=unused-import
+    exog_ss_mod,
+    make_statespace_mod,
+    ss_mod,
+)
+from pymc_experimental.tests.statespace.utilities.shared_fixtures import (  # pylint: disable=unused-import
+    rng,
+)
+from pymc_experimental.tests.statespace.utilities.test_helpers import (
+    load_nile_test_data,
+)
+
+pytest.importorskip("jax")
+
+
+floatX = pytensor.config.floatX
+nile = load_nile_test_data()
+ALL_SAMPLE_OUTPUTS = MATRIX_NAMES + FILTER_OUTPUT_NAMES + SMOOTHER_OUTPUT_NAMES
+
+
+@pytest.fixture(scope="session")
+def pymc_mod(ss_mod):
+    with pm.Model(coords=ss_mod.coords) as pymc_mod:
+        rho = pm.Beta("rho", 1, 1)
+        zeta = pm.Deterministic("zeta", 1 - rho)
+
+        ss_mod.build_statespace_graph(
+            data=nile, mode="JAX", save_kalman_filter_outputs_in_idata=True
+        )
+        names = ["x0", "P0", "c", "d", "T", "Z", "R", "H", "Q"]
+        for name, matrix in zip(names, ss_mod.unpack_statespace()):
+            pm.Deterministic(name, matrix)
+
+    return pymc_mod
+
+
+@pytest.fixture(scope="session")
+def exog_pymc_mod(exog_ss_mod, rng):
+    y = rng.normal(size=(100, 1)).astype(floatX)
+    X = rng.normal(size=(100, 3)).astype(floatX)
+
+    with pm.Model(coords=exog_ss_mod.coords) as m:
+        exog_data = pm.Data("data_exog", X)
+        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
+        P0_sigma = pm.Exponential("P0_sigma", 1)
+        P0 = pm.Deterministic(
+            "P0", pt.eye(exog_ss_mod.k_states) * P0_sigma, dims=["state", "state_aux"]
+        )
+        beta_exog = pm.Normal("beta_exog", dims=["exog_state"])
+
+        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
+        exog_ss_mod.build_statespace_graph(y, mode="JAX")
+
+    return m
+
+
+@pytest.fixture(scope="session")
+def idata(pymc_mod, rng):
+    with warnings.catch_warnings(action="ignore"):
+        with pymc_mod:
+            idata = pm.sample(
+                draws=10,
+                tune=1,
+                chains=1,
+                random_seed=rng,
+                nuts_sampler="numpyro",
+                progressbar=False,
+            )
+        with freeze_dims_and_data(pymc_mod):
+            idata_prior = pm.sample_prior_predictive(
+                samples=10, random_seed=rng, compile_kwargs={"mode": "JAX"}
+            )
+
+    idata.extend(idata_prior)
+    return idata
+
+
+@pytest.fixture(scope="session")
+def idata_exog(exog_pymc_mod, rng):
+    with warnings.catch_warnings(action="ignore"):
+        with exog_pymc_mod:
+            idata = pm.sample(
+                draws=10,
+                tune=1,
+                chains=1,
+                random_seed=rng,
+                nuts_sampler="numpyro",
+                progressbar=False,
+            )
+        with freeze_dims_and_data(pymc_mod):
+            idata_prior = pm.sample_prior_predictive(
+                samples=10, random_seed=rng, compile_kwargs={"mode": "JAX"}
+            )
+
+    idata.extend(idata_prior)
+    return idata
+
+
+@pytest.mark.parametrize("group", ["posterior", "prior"])
+@pytest.mark.parametrize("matrix", ALL_SAMPLE_OUTPUTS)
+def test_no_nans_in_sampling_output(ss_mod, group, matrix, idata):
+    assert not np.any(np.isnan(idata[group][matrix].values))
+
+
+@pytest.mark.parametrize("group", ["prior", "posterior"])
+@pytest.mark.parametrize("kind", ["conditional", "unconditional"])
+def test_sampling_methods(group, kind, ss_mod, idata, rng):
+    assert ss_mod._fit_mode == "JAX"
+
+    f = getattr(ss_mod, f"sample_{kind}_{group}")
+    with pytest.warns(UserWarning, match="The RandomType SharedVariables"):
+        test_idata = f(idata, random_seed=rng)
+
+    if kind == "conditional":
+        for output in ["filtered", "predicted", "smoothed"]:
+            assert f"{output}_{group}" in test_idata
+            assert not np.any(np.isnan(test_idata[f"{output}_{group}"].values))
+            assert not np.any(np.isnan(test_idata[f"{output}_{group}_observed"].values))
+
+    if kind == "unconditional":
+        for output in ["latent", "observed"]:
+            assert f"{group}_{output}" in test_idata
+            assert not np.any(np.isnan(test_idata[f"{group}_{output}"].values))
+
+
+@pytest.mark.parametrize("filter_output", ["predicted", "filtered", "smoothed"])
+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
+        )
+
+    assert forecast_idata.coords["time"].values.shape == (10,)
+    assert forecast_idata.forecast_latent.dims == ("chain", "draw", "time", "state")
+    assert forecast_idata.forecast_observed.dims == ("chain", "draw", "time", "observed_state")
+
+    assert not np.any(np.isnan(forecast_idata.forecast_latent.values))
+    assert not np.any(np.isnan(forecast_idata.forecast_observed.values))