Remove mode argument from Statespace models

pymc_extras/statespace/core/statespace.py

@@ -14,7 +14,6 @@
 from pymc.model.transform.optimization import freeze_dims_and_data
 from pymc.util import RandomState
 from pytensor import Variable, graph_replace
-from pytensor.compile import get_mode
 from rich.box import SIMPLE_HEAD
 from rich.console import Console
 from rich.table import Table
@@ -222,7 +221,6 @@ def __init__(
         verbose: bool = True,
         measurement_error: bool = False,
     ):
-        self._fit_mode: str | None = None
         self._fit_coords: dict[str, Sequence[str]] | None = None
         self._fit_dims: dict[str, Sequence[str]] | None = None
         self._fit_data: pt.TensorVariable | None = None
@@ -819,7 +817,6 @@ def build_statespace_graph(
         self,
         data: np.ndarray | pd.DataFrame | pt.TensorVariable,
         register_data: bool = True,
-        mode: str | None = None,
         missing_fill_value: float | None = None,
         cov_jitter: float | None = JITTER_DEFAULT,
         save_kalman_filter_outputs_in_idata: bool = False,
@@ -889,7 +886,6 @@ def build_statespace_graph(
         filter_outputs = self.kalman_filter.build_graph(
             pt.as_tensor_variable(data),
             *self.unpack_statespace(),
-            mode=mode,
             missing_fill_value=missing_fill_value,
             cov_jitter=cov_jitter,
         )
@@ -900,7 +896,7 @@ def build_statespace_graph(
         filtered_covariances, predicted_covariances, observed_covariances = covs
         if save_kalman_filter_outputs_in_idata:
             smooth_states, smooth_covariances = self._build_smoother_graph(
-                filtered_states, filtered_covariances, self.unpack_statespace(), mode=mode
+                filtered_states, filtered_covariances, self.unpack_statespace()
             )
             all_kf_outputs = [*states, smooth_states, *covs, smooth_covariances]
             self._register_kalman_filter_outputs_with_pymc_model(all_kf_outputs)
@@ -919,7 +915,6 @@ def build_statespace_graph(
 
         self._fit_coords = pm_mod.coords.copy()
         self._fit_dims = pm_mod.named_vars_to_dims.copy()
-        self._fit_mode = mode
 
     def _build_smoother_graph(
         self,
@@ -964,7 +959,7 @@ def _build_smoother_graph(
             *_, T, Z, R, H, Q = matrices
 
             smooth_states, smooth_covariances = self.kalman_smoother.build_graph(
-                T, R, Q, filtered_states, filtered_covariances, mode=mode, cov_jitter=cov_jitter
+                T, R, Q, filtered_states, filtered_covariances, cov_jitter=cov_jitter
             )
             smooth_states.name = "smooth_states"
             smooth_covariances.name = "smooth_covariances"
@@ -1082,7 +1077,6 @@ def _kalman_filter_outputs_from_dummy_graph(
             R,
             H,
             Q,
-            mode=self._fit_mode,
         )
 
         filter_outputs.pop(-1)
@@ -1092,7 +1086,7 @@ def _kalman_filter_outputs_from_dummy_graph(
         filtered_covariances, predicted_covariances, _ = covariances
 
         [smoothed_states, smoothed_covariances] = self.kalman_smoother.build_graph(
-            T, R, Q, filtered_states, filtered_covariances, mode=self._fit_mode
+            T, R, Q, filtered_states, filtered_covariances
         )
 
         grouped_outputs = [
@@ -1208,7 +1202,6 @@ def _sample_conditional(
                     for name in FILTER_OUTPUT_TYPES
                     for suffix in ["", "_observed"]
                 ],
-                compile_kwargs={"mode": get_mode(self._fit_mode)},
                 random_seed=random_seed,
                 **kwargs,
             )
@@ -1308,7 +1301,6 @@ def _sample_unconditional(
                 *matrices,
                 steps=steps,
                 dims=dims,
-                mode=self._fit_mode,
                 sequence_names=self.kalman_filter.seq_names,
                 k_endog=self.k_endog,
             )
@@ -1323,7 +1315,6 @@ def _sample_unconditional(
             idata_unconditional = pm.sample_posterior_predictive(
                 group_idata,
                 var_names=[f"{group}_latent", f"{group}_observed"],
-                compile_kwargs={"mode": self._fit_mode},
                 random_seed=random_seed,
                 **kwargs,
             )
@@ -1547,7 +1538,6 @@ def sample_statespace_matrices(
             matrix_idata = pm.sample_posterior_predictive(
                 idata if group == "posterior" else idata.prior,
                 var_names=matrix_names,
-                compile_kwargs={"mode": self._fit_mode},
                 extend_inferencedata=False,
             )
 
@@ -2094,7 +2084,6 @@ def forecast(
                 *matrices,
                 steps=len(forecast_index),
                 dims=dims,
-                mode=self._fit_mode,
                 sequence_names=self.kalman_filter.seq_names,
                 k_endog=self.k_endog,
                 append_x0=False,
@@ -2109,7 +2098,6 @@ def forecast(
             idata_forecast = pm.sample_posterior_predictive(
                 idata,
                 var_names=["forecast_latent", "forecast_observed"],
-                compile_kwargs={"mode": self._fit_mode},
                 random_seed=random_seed,
                 **kwargs,
             )
@@ -2260,28 +2248,13 @@ def irf_step(shock, x, c, T, R):
                 non_sequences=[c, T, R],
                 n_steps=n_steps,
                 strict=True,
-                mode=self._fit_mode,
             )
 
             pm.Deterministic("irf", irf, dims=[TIME_DIM, ALL_STATE_DIM])
 
-            compile_kwargs = kwargs.get("compile_kwargs", {})
-            if "mode" not in compile_kwargs.keys():
-                compile_kwargs = {"mode": self._fit_mode}
-            else:
-                mode = compile_kwargs.get("mode")
-                if mode is not None and mode != self._fit_mode:
-                    raise ValueError(
-                        f"User provided compile mode ({mode}) does not match the compile mode used to "
-                        f"construct the model ({self._fit_mode})."
-                    )
-
-                compile_kwargs.update({"mode": self._fit_mode})
-
             irf_idata = pm.sample_posterior_predictive(
                 idata,
                 var_names=["irf"],
-                compile_kwargs=compile_kwargs,
                 random_seed=random_seed,
                 **kwargs,
             )

pymc_extras/statespace/filters/distributions.py

@@ -69,7 +69,6 @@ def __new__(
         H,
         Q,
         steps=None,
-        mode=None,
         sequence_names=None,
         append_x0=True,
         **kwargs,
@@ -97,7 +96,6 @@ def __new__(
             H,
             Q,
             steps=steps,
-            mode=mode,
             sequence_names=sequence_names,
             append_x0=append_x0,
             **kwargs,
@@ -116,7 +114,6 @@ def dist(
         H,
         Q,
         steps=None,
-        mode=None,
         sequence_names=None,
         append_x0=True,
         **kwargs,
@@ -132,7 +129,6 @@ def dist(
 
         return super().dist(
             [a0, P0, c, d, T, Z, R, H, Q, steps],
-            mode=mode,
             sequence_names=sequence_names,
             append_x0=append_x0,
             **kwargs,
@@ -152,7 +148,6 @@ def rv_op(
         Q,
         steps,
         size=None,
-        mode=None,
         sequence_names=None,
         append_x0=True,
     ):
@@ -235,7 +230,6 @@ def step_fn(*args):
             sequences=None if len(sequences) == 0 else sequences,
             non_sequences=[*non_sequences, rng],
             n_steps=steps,
-            mode=mode,
             strict=True,
         )
 
@@ -279,7 +273,6 @@ def __new__(
         steps,
         k_endog=None,
         sequence_names=None,
-        mode=None,
         append_x0=True,
         **kwargs,
     ):
@@ -307,7 +300,6 @@ def __new__(
             H,
             Q,
             steps=steps,
-            mode=mode,
             sequence_names=sequence_names,
             append_x0=append_x0,
             **kwargs,

pymc_extras/statespace/filters/kalman_filter.py

@@ -5,7 +5,6 @@
 import pytensor.tensor as pt
 
 from pymc.pytensorf import constant_fold
-from pytensor.compile.mode import get_mode
 from pytensor.graph.basic import Variable
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
@@ -28,25 +27,17 @@
 
 
 class BaseFilter(ABC):
-    def __init__(self, mode=None):
+    def __init__(self):
         """
         Kalman Filter.
 
-        Parameters
-        ----------
-        mode : str, optional
-            The mode used for Pytensor compilation. Defaults to None.
-
         Notes
         -----
         The BaseFilter class is an abstract base class (ABC) for implementing kalman filters.
         It defines common attributes and methods used by kalman filter implementations.
 
         Attributes
         ----------
-        mode : str or None
-            The mode used for Pytensor compilation.
-
         seq_names : list[str]
             A list of name representing time-varying statespace matrices. That is, inputs that will need to be
             provided to the `sequences` argument of `pytensor.scan`
@@ -56,7 +47,6 @@ def __init__(self, mode=None):
             to the `non_sequences` argument of `pytensor.scan`
         """
 
-        self.mode: str = mode
         self.seq_names: list[str] = []
         self.non_seq_names: list[str] = []
 
@@ -153,7 +143,6 @@ def build_graph(
         R,
         H,
         Q,
-        mode=None,
         return_updates=False,
         missing_fill_value=None,
         cov_jitter=None,
@@ -166,9 +155,6 @@ def build_graph(
         data : TensorVariable
             Data to be filtered
 
-        mode : optional, str
-            Pytensor compile mode, passed to pytensor.scan
-
         return_updates: bool, default False
             Whether to return updates associated with the pytensor scan. Should only be requried to debug pruposes.
 
@@ -199,7 +185,6 @@ def build_graph(
         if cov_jitter is None:
             cov_jitter = JITTER_DEFAULT
 
-        self.mode = mode
         self.missing_fill_value = missing_fill_value
         self.cov_jitter = cov_jitter
 
@@ -227,7 +212,6 @@ def build_graph(
             outputs_info=[None, a0, None, None, P0, None, None],
             non_sequences=non_sequences,
             name="forward_kalman_pass",
-            mode=get_mode(self.mode),
             strict=False,
         )
 
@@ -800,7 +784,6 @@ def kalman_step(self, y, a, P, c, d, T, Z, R, H, Q):
             self._univariate_inner_filter_step,
             sequences=[y_masked, Z_masked, d, pt.diag(H_masked), nan_mask],
             outputs_info=[a, P, None, None, None],
-            mode=get_mode(self.mode),
             name="univariate_inner_scan",
         )
 

pymc_extras/statespace/filters/kalman_smoother.py

@@ -1,7 +1,6 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.compile import get_mode
 from pytensor.tensor.nlinalg import matrix_dot
 
 from pymc_extras.statespace.filters.utilities import (
@@ -18,8 +17,7 @@ class KalmanSmoother:
 
     """
 
-    def __init__(self, mode: str | None = None):
-        self.mode = mode
+    def __init__(self):
         self.cov_jitter = JITTER_DEFAULT
         self.seq_names = []
         self.non_seq_names = []
@@ -64,9 +62,8 @@ def unpack_args(self, args):
         return a, P, a_smooth, P_smooth, T, R, Q
 
     def build_graph(
-        self, T, R, Q, filtered_states, filtered_covariances, mode=None, cov_jitter=JITTER_DEFAULT
+        self, T, R, Q, filtered_states, filtered_covariances, cov_jitter=JITTER_DEFAULT
     ):
-        self.mode = mode
         self.cov_jitter = cov_jitter
 
         n, k = filtered_states.type.shape
@@ -88,7 +85,6 @@ def build_graph(
             non_sequences=non_sequences,
             go_backwards=True,
             name="kalman_smoother",
-            mode=get_mode(self.mode),
         )
 
         smoothed_states, smoothed_covariances = smoother_result

pymc_extras/statespace/models/SARIMAX.py

@@ -158,7 +158,7 @@ class BayesianSARIMA(PyMCStateSpace):
             rho = pm.Beta("ar_params", alpha=5, beta=1, dims=ss_mod.param_dims["ar_params"])
             theta = pm.Normal("ma_params", mu=0.0, sigma=0.5, dims=ss_mod.param_dims["ma_params"])
 
-            ss_mod.build_statespace_graph(df, mode="JAX")
+            ss_mod.build_statespace_graph(df)
             idata = pm.sample(nuts_sampler='numpyro')
 
     References
@@ -366,17 +366,15 @@ def coords(self) -> dict[str, Sequence]:
 
         return coords
 
-    def _stationary_initialization(self, mode=None):
+    def _stationary_initialization(self):
         # Solve for matrix quadratic for P0
         T = self.ssm["transition"]
         R = self.ssm["selection"]
         Q = self.ssm["state_cov"]
         c = self.ssm["state_intercept"]
 
         x0 = pt.linalg.solve(pt.identity_like(T) - T, c, assume_a="gen", check_finite=True)
-
-        method = "direct" if (self.k_states < 5) or (mode == "JAX") else "bilinear"
-        P0 = solve_discrete_lyapunov(T, pt.linalg.matrix_dot(R, Q, R.T), method=method)
+        P0 = solve_discrete_lyapunov(T, pt.linalg.matrix_dot(R, Q, R.T), method="bilinear")
 
         return x0, P0
 

pymc_extras/statespace/models/VARMAX.py

@@ -135,7 +135,7 @@ class BayesianVARMAX(PyMCStateSpace):
             ar_params = pm.Normal("ar_params", mu=0, sigma=1, dims=ar_dims)
             state_cov = pm.Deterministic("state_cov", state_chol @ state_chol.T, dims=state_cov_dims)
 
-            bvar_mod.build_statespace_graph(data, mode="JAX")
+            bvar_mod.build_statespace_graph(data)
             idata = pm.sample(nuts_sampler="numpyro")
     """