Add deprecation warning to mode argument of build_statespace_graph

Add mode argument to statespace constructors

Author: jessegrabowski

pymc_extras/statespace/core/statespace.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 
 from collections.abc import Callable, Sequence
 from typing import Any, Literal
@@ -98,6 +99,13 @@ class PyMCStateSpace:
         compute the observation errors. If False, these errors are deterministically zero; if True, they are sampled
         from a multivariate normal.
 
+    mode: str or Mode, optional
+        Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+        ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+        Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+        to all sampling methods.
+
     Notes
     -----
     Based on the statsmodels statespace implementation https://github.com/statsmodels/statsmodels/blob/main/statsmodels/tsa/statespace/representation.py,
@@ -220,6 +228,7 @@ def __init__(
         filter_type: str = "standard",
         verbose: bool = True,
         measurement_error: bool = False,
+        mode: str | None = None,
     ):
         self._fit_coords: dict[str, Sequence[str]] | None = None
         self._fit_dims: dict[str, Sequence[str]] | None = None
@@ -235,6 +244,7 @@ def __init__(
         self.k_states = k_states
         self.k_posdef = k_posdef
         self.measurement_error = measurement_error
+        self.mode = mode
 
         # All models contain a state space representation and a Kalman filter
         self.ssm = PytensorRepresentation(k_endog, k_states, k_posdef)
@@ -821,6 +831,7 @@ def build_statespace_graph(
         cov_jitter: float | None = JITTER_DEFAULT,
         mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         save_kalman_filter_outputs_in_idata: bool = False,
+        mode: str | None = None,
     ) -> None:
         """
         Given a parameter vector `theta`, constructs the full computational graph describing the state space model and
@@ -874,7 +885,25 @@ def build_statespace_graph(
         save_kalman_filter_outputs_in_idata: bool, optional, default=False
             If True, Kalman Filter outputs will be saved in the model as deterministics. Useful for debugging, but
             should not be necessary for the majority of users.
+
+        mode: str, optional
+            Pytensor mode to use when compiling the graph. This will be saved as a model attribute and used when
+            compiling sampling functions (e.g. ``sample_conditional_prior``).
+
+            .. deprecated:: 0.2.5
+                The `mode` argument is deprecated and will be removed in a future version. Pass ``mode`` to the
+                model constructor, or manually specify ``compile_kwargs`` in sampling functions instead.
+
         """
+        if mode is not None:
+            warnings.warn(
+                "The `mode` argument is deprecated and will be removed in a future version. "
+                "Pass `mode` to the model constructor, or manually specify `compile_kwargs` in sampling functions"
+                " instead.",
+                DeprecationWarning,
+            )
+            self.mode = mode
+
         pm_mod = modelcontext(None)
 
         self._insert_random_variables()
@@ -1107,6 +1136,12 @@ def _kalman_filter_outputs_from_dummy_graph(
 
         return [x0, P0, c, d, T, Z, R, H, Q], grouped_outputs
 
+    def _set_default_mode(self, compile_kwargs):
+        mode = compile_kwargs.get("mode", self.mode)
+        compile_kwargs["mode"] = mode
+
+        return compile_kwargs
+
     def _sample_conditional(
         self,
         idata: InferenceData,
@@ -1158,6 +1193,9 @@ def _sample_conditional(
         _verify_group(group)
         group_idata = getattr(idata, group)
 
+        compile_kwargs = kwargs.pop("compile_kwargs", {})
+        compile_kwargs = self._set_default_mode(compile_kwargs)
+
         with pm.Model(coords=self._fit_coords) as forward_model:
             (
                 [
@@ -1224,6 +1262,7 @@ def _sample_conditional(
                     for suffix in ["", "_observed"]
                 ],
                 random_seed=random_seed,
+                compile_kwargs=compile_kwargs,
                 **kwargs,
             )
 
@@ -1289,6 +1328,10 @@ def _sample_unconditional(
               the latent state trajectories: `y[t] = Z @ x[t] + nu[t]`, where `nu ~ N(0, H)`.
         """
         _verify_group(group)
+
+        compile_kwargs = kwargs.pop("compile_kwargs", {})
+        compile_kwargs = self._set_default_mode(compile_kwargs)
+
         group_idata = getattr(idata, group)
         dims = None
         temp_coords = self._fit_coords.copy()
@@ -1347,6 +1390,7 @@ def _sample_unconditional(
                 group_idata,
                 var_names=[f"{group}_latent", f"{group}_observed"],
                 random_seed=random_seed,
+                compile_kwargs=compile_kwargs,
                 **kwargs,
             )
 
@@ -1574,7 +1618,7 @@ def sample_unconditional_posterior(
         )
 
     def sample_statespace_matrices(
-        self, idata, matrix_names: str | list[str] | None, group: str = "posterior"
+        self, idata, matrix_names: str | list[str] | None, group: str = "posterior", **kwargs
     ):
         """
         Draw samples of requested statespace matrices from provided idata
@@ -1591,12 +1635,18 @@ def sample_statespace_matrices(
         group: str, one of "posterior" or "prior"
             Whether to sample from priors or posteriors
 
+        kwargs:
+            Additional keyword arguments are passed to ``pymc.sample_posterior_predictive``
+
         Returns
         -------
         idata_matrices: az.InterenceData
         """
         _verify_group(group)
 
+        compile_kwargs = kwargs.pop("compile_kwargs", {})
+        compile_kwargs = self._set_default_mode(compile_kwargs)
+
         if matrix_names is None:
             matrix_names = MATRIX_NAMES
         elif isinstance(matrix_names, str):
@@ -1628,6 +1678,8 @@ def sample_statespace_matrices(
                 idata if group == "posterior" else idata.prior,
                 var_names=matrix_names,
                 extend_inferencedata=False,
+                compile_kwargs=compile_kwargs,
+                **kwargs,
             )
 
         return matrix_idata
@@ -2096,6 +2148,10 @@ def forecast(
         filter_time_dim = TIME_DIM
 
         _validate_filter_arg(filter_output)
+
+        compile_kwargs = kwargs.pop("compile_kwargs", {})
+        compile_kwargs = self._set_default_mode(compile_kwargs)
+
         time_index = self._get_fit_time_index()
 
         if start is None and verbose:
@@ -2198,6 +2254,7 @@ def forecast(
                 idata,
                 var_names=["forecast_latent", "forecast_observed"],
                 random_seed=random_seed,
+                compile_kwargs=compile_kwargs,
                 **kwargs,
             )
 
@@ -2285,6 +2342,9 @@ def impulse_response_function(
         n_options = sum(x is not None for x in options)
         Q = None  # No covariance matrix needed if a trajectory is provided. Will be overwritten later if needed.
 
+        compile_kwargs = kwargs.pop("compile_kwargs", {})
+        compile_kwargs = self._set_default_mode(compile_kwargs)
+
         if n_options > 1:
             raise ValueError("Specify exactly 0 or 1 of shock_size, shock_cov, or shock_trajectory")
         elif n_options == 1:
@@ -2364,6 +2424,7 @@ def irf_step(shock, x, c, T, R):
                 idata,
                 var_names=["irf"],
                 random_seed=random_seed,
+                compile_kwargs=compile_kwargs,
                 **kwargs,
             )
 

pymc_extras/statespace/models/ETS.py

@@ -5,6 +5,7 @@
 import pytensor.tensor as pt
 
 from pytensor import graph_replace
+from pytensor.compile.mode import Mode
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
@@ -35,6 +36,7 @@ def __init__(
         initialization_dampening: float = 0.8,
         filter_type: str = "standard",
         verbose: bool = True,
+        mode: str | Mode | None = None,
     ):
         r"""
         Exponential Smoothing State Space Model
@@ -212,6 +214,13 @@ def __init__(
             and "cholesky". See the docs for kalman filters for more details.
         verbose: bool, default True
             If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
+        mode: str or Mode, optional
+            Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+            ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+            Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+            to all sampling methods.
+
 
         References
         ----------
@@ -284,6 +293,7 @@ def __init__(
             filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
 
     @property

pymc_extras/statespace/models/SARIMAX.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pytensor.tensor as pt
 
+from pytensor.compile.mode import Mode
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace, floatX
@@ -91,6 +92,13 @@ class BayesianSARIMA(PyMCStateSpace):
     verbose: bool, default True
         If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
 
+    mode: str or Mode, optional
+        Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+        ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+        Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+        to all sampling methods.
+
     Notes
     -----
         The ARIMAX model is a univariate time series model that posits the future evolution of a stationary time series will
@@ -180,7 +188,21 @@ def __init__(
         state_structure: str = "fast",
         measurement_error: bool = False,
         verbose=True,
+        mode: str | Mode | None = None,
     ):
+        """
+
+        Parameters
+        ----------
+        order
+        seasonal_order
+        stationary_initialization
+        filter_type
+        state_structure
+        measurement_error
+        verbose
+        mode
+        """
         # Model order
         self.p, self.d, self.q = order
         if seasonal_order is None:
@@ -228,6 +250,7 @@ def __init__(
             filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
 
     @property

pymc_extras/statespace/models/VARMAX.py

@@ -5,6 +5,7 @@
 import pytensor
 import pytensor.tensor as pt
 
+from pytensor.compile.mode import Mode
 from pytensor.tensor.slinalg import solve_discrete_lyapunov
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
@@ -72,6 +73,13 @@ class BayesianVARMAX(PyMCStateSpace):
     verbose: bool, default True
         If true, a message will be logged to the terminal explaining the variable names, dimensions, and supports.
 
+    mode: str or Mode, optional
+        Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+        ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+        Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+        to all sampling methods.
+
     Notes
     -----
     The VARMA model is a multivariate extension of the SARIMAX model. Given a set of timeseries :math:`\{x_t\}_{t=0}^T`,
@@ -147,7 +155,8 @@ def __init__(
         stationary_initialization: bool = False,
         filter_type: str = "standard",
         measurement_error: bool = False,
-        verbose=True,
+        verbose: bool = True,
+        mode: str | Mode | None = None,
     ):
         if (endog_names is None) and (k_endog is None):
             raise ValueError("Must specify either endog_names or k_endog")
@@ -174,6 +183,7 @@ def __init__(
             filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
 
         # Save counts of the number of parameters in each category

pymc_extras/statespace/models/structural.py

@@ -12,6 +12,7 @@
 import xarray as xr
 
 from pytensor import Variable
+from pytensor.compile.mode import Mode
 
 from pymc_extras.statespace.core import PytensorRepresentation
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
@@ -81,6 +82,7 @@ def __init__(
         name: str | None = None,
         verbose: bool = True,
         filter_type: str = "standard",
+        mode: str | Mode | None = None,
     ):
         # Add the initial state covariance to the parameters
         if name is None:
@@ -112,6 +114,7 @@ def __init__(
             filter_type=filter_type,
             verbose=verbose,
             measurement_error=measurement_error,
+            mode=mode,
         )
         self.ssm = ssm.copy()
 
@@ -644,7 +647,9 @@ def __add__(self, other):
 
         return new_comp
 
-    def build(self, name=None, filter_type="standard", verbose=True):
+    def build(
+        self, name=None, filter_type="standard", verbose=True, mode: str | Mode | None = None
+    ):
         """
         Build a StructuralTimeSeries statespace model from the current component(s)
 
@@ -660,6 +665,13 @@ def build(self, name=None, filter_type="standard", verbose=True):
         verbose : bool, optional
             If True, displays information about the initialized model. Defaults to True.
 
+        mode: str or Mode, optional
+            Pytensor compile mode, used in auxiliary sampling methods such as ``sample_conditional_posterior`` and
+            ``forecast``. The mode does **not** effect calls to ``pm.sample``.
+
+            Regardless of whether a mode is specified, it can always be overwritten via the ``compile_kwargs`` argument
+            to all sampling methods.
+
         Returns
         -------
         PyMCStateSpace
@@ -685,6 +697,7 @@ def build(self, name=None, filter_type="standard", verbose=True):
             name_to_data=self._name_to_data,
             filter_type=filter_type,
             verbose=verbose,
+            mode=mode,
         )
 
 

tests/statespace/models/test_ETS.py

@@ -89,6 +89,12 @@ def test_order_flags(order, expected_flags):
         assert getattr(mod, key) == value
 
 
+def tests_mode_argument():
+    # Mode argument should be passed to the parent class
+    mod = BayesianETS(order=("A", "N", "N"), mode="FAST_RUN")
+    assert mod.mode == "FAST_RUN"
+
+
 @pytest.mark.parametrize("order, expected_params", zip(orders, order_params), ids=order_names)
 def test_param_info(order: tuple[str, str, str], expected_params):
     mod = BayesianETS(order=order, seasonal_periods=4)