Working with Filter, not with Smoother

Author: aandorra-mia

notebooks/batch-examples.ipynb

@@ -189,7 +189,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 11,
    "id": "1262c7d4",
    "metadata": {},
    "outputs": [],
@@ -199,7 +199,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 12,
    "id": "2dcd3958",
    "metadata": {},
    "outputs": [
@@ -228,7 +228,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 13,
    "id": "6f41344f",
    "metadata": {},
    "outputs": [],
@@ -238,7 +238,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 14,
    "id": "44905b8a",
    "metadata": {},
    "outputs": [],
@@ -248,7 +248,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 15,
    "id": "34fe01b8",
    "metadata": {},
    "outputs": [
@@ -258,7 +258,7 @@
        "(3, 10, 5)"
       ]
      },
-     "execution_count": 24,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -270,7 +270,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 16,
    "id": "f37efe79",
    "metadata": {},
    "outputs": [
@@ -288,7 +288,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 17,
    "id": "7b45de74",
    "metadata": {},
    "outputs": [
@@ -298,7 +298,7 @@
        "(3, 10)"
       ]
      },
-     "execution_count": 26,
+     "execution_count": 17,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -317,7 +317,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 18,
    "id": "69519822",
    "metadata": {},
    "outputs": [],
@@ -327,17 +327,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 19,
    "id": "3f745449",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "633 μs ± 18.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
-      "1.52 ms ± 35.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
-      "4.76 ms ± 259 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+      "675 μs ± 22.3 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
+      "1.64 ms ± 37.5 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
+      "5.28 ms ± 424 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
      ]
     }
    ],

pymc_extras/statespace/filters/kalman_filter.py

@@ -10,6 +10,7 @@
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.slinalg import solve_triangular
+from pytensor.graph.replace import vectorize_graph
 
 from pymc_extras.statespace.filters.utilities import (
     quad_form_sym,
@@ -20,6 +21,7 @@
 
 MVN_CONST = pt.log(2 * pt.constant(np.pi, dtype="float64"))
 PARAM_NAMES = ["c", "d", "T", "Z", "R", "H", "Q"]
+CORE_NDIM = (2, 1, 2, 1, 1, 2, 2, 2, 2, 2)
 
 assert_time_varying_dim_correct = Assert(
     "The first dimension of a time varying matrix (the time dimension) must be "
@@ -73,6 +75,23 @@ def check_params(self, data, a0, P0, c, d, T, Z, R, H, Q):
         """
         return data, a0, P0, c, d, T, Z, R, H, Q
 
+    def has_batched_input(self, data, a0, P0, c, d, T, Z, R, H, Q):
+        """
+        Check if any of the inputs are batched.
+        """
+        return any(x.ndim > CORE_NDIM[i] for i, x in enumerate([data, a0, P0, c, d, T, Z, R, H, Q]))
+
+    def get_dummy_core_inputs(self, data, a0, P0, c, d, T, Z, R, H, Q):
+        """
+        Get dummy inputs for the core parameters.
+        """
+        out = []
+        for x, core_ndim in zip([data, a0, P0, c, d, T, Z, R, H, Q], CORE_NDIM):
+            out.append(
+                pt.tensor(f"{x.name}_core_case", dtype=x.dtype, shape=x.type.shape[-core_ndim:])
+            )
+        return out
+
     @staticmethod
     def add_check_on_time_varying_shapes(
         data: TensorVariable, sequence_params: list[TensorVariable]
@@ -202,13 +221,18 @@ def build_graph(
         self.mode = mode
         self.missing_fill_value = missing_fill_value
         self.cov_jitter = cov_jitter
+        is_batched = self.has_batched_input(data, a0, P0, c, d, T, Z, R, H, Q)
 
         [R_shape] = constant_fold([R.shape], raise_not_constant=False)
         [Z_shape] = constant_fold([Z.shape], raise_not_constant=False)
 
         self.n_states, self.n_shocks = R_shape[-2:]
         self.n_endog = Z_shape[-2]
 
+        if is_batched:
+            batched_inputs = [data, a0, P0, c, d, T, Z, R, H, Q]
+            data, a0, P0, c, d, T, Z, R, H, Q = self.get_dummy_core_inputs(*batched_inputs)
+
         data, a0, P0, *params = self.check_params(data, a0, P0, c, d, T, Z, R, H, Q)
 
         sequences, non_sequences, seq_names, non_seq_names = split_vars_into_seq_and_nonseq(
@@ -233,8 +257,13 @@ def build_graph(
 
         filter_results = self._postprocess_scan_results(results, a0, P0, n=data.type.shape[0])
 
+        if is_batched:
+            vec_subs = dict(zip([data, a0, P0, c, d, T, Z, R, H, Q], batched_inputs))
+            filter_results = vectorize_graph(filter_results, vec_subs)
+
         if return_updates:
             return filter_results, updates
+
         return filter_results
 
     def _postprocess_scan_results(self, results, a0, P0, n) -> list[TensorVariable]:

pymc_extras/statespace/filters/kalman_smoother.py

@@ -3,14 +3,16 @@
 
 from pytensor.compile import get_mode
 from pytensor.tensor.nlinalg import matrix_dot
-
+from pytensor.graph.replace import vectorize_graph
 from pymc_extras.statespace.filters.utilities import (
     quad_form_sym,
     split_vars_into_seq_and_nonseq,
     stabilize,
 )
 from pymc_extras.statespace.utils.constants import JITTER_DEFAULT
 
+SMOOTHER_CORE_NDIM = (2, 2, 2, 2, 3)
+
 
 class KalmanSmoother:
     """
@@ -63,12 +65,41 @@ def unpack_args(self, args):
 
         return a, P, a_smooth, P_smooth, T, R, Q
 
+    def has_batched_input(self, T, R, Q, filtered_states, filtered_covariances):
+        """
+        Check if any of the inputs are batched.
+        """
+        return any(
+            x.ndim > SMOOTHER_CORE_NDIM[i]
+            for i, x in enumerate([T, R, Q, filtered_states, filtered_covariances])
+        )
+
+    def get_dummy_core_inputs(self, T, R, Q, filtered_states, filtered_covariances):
+        """
+        Get dummy inputs for the core parameters.
+        """
+        out = []
+        for x, core_ndim in zip(
+            [T, R, Q, filtered_states, filtered_covariances], SMOOTHER_CORE_NDIM
+        ):
+            out.append(
+                pt.tensor(f"{x.name}_core_case", dtype=x.dtype, shape=x.type.shape[-core_ndim:])
+            )
+        return out
+
     def build_graph(
         self, T, R, Q, filtered_states, filtered_covariances, mode=None, cov_jitter=JITTER_DEFAULT
     ):
         self.mode = mode
         self.cov_jitter = cov_jitter
 
+        is_batched = self.has_batched_input(T, R, Q, filtered_states, filtered_covariances)
+        if is_batched:
+            batched_inputs = [T, R, Q, filtered_states, filtered_covariances]
+            T, R, Q, filtered_states, filtered_covariances = self.get_dummy_core_inputs(
+                *batched_inputs
+            )
+
         n, k = filtered_states.type.shape
 
         a_last = pt.specify_shape(filtered_states[-1], (k,))
@@ -98,6 +129,12 @@ def build_graph(
         smoothed_covariances = pt.concatenate(
             [smoothed_covariances[::-1], pt.expand_dims(P_last, axis=(0,))], axis=0
         )
+        smoothed_states.dprint()
+        if is_batched:
+            vec_subs = dict(zip([T, R, Q, filtered_states, filtered_covariances], batched_inputs))
+            smoothed_states, smoothed_covariances = vectorize_graph(
+                [smoothed_states, smoothed_covariances], vec_subs
+            )
 
         smoothed_states.name = "smoothed_states"
         smoothed_covariances.name = "smoothed_covariances"

tests/statespace/test_kalman_filter.py

@@ -31,19 +31,22 @@
 RTOL = 1e-6 if floatX.endswith("64") else 1e-3
 
 standard_inout = initialize_filter(StandardFilter())
+standard_inout_batched = initialize_filter(StandardFilter(), batched=True)
 cholesky_inout = initialize_filter(SquareRootFilter())
 univariate_inout = initialize_filter(UnivariateFilter())
 
 f_standard = pytensor.function(*standard_inout, on_unused_input="ignore")
+f_standard_batched = pytensor.function(*standard_inout_batched, on_unused_input="ignore")
 f_cholesky = pytensor.function(*cholesky_inout, on_unused_input="ignore")
 f_univariate = pytensor.function(*univariate_inout, on_unused_input="ignore")
 
-filter_funcs = [f_standard, f_cholesky, f_univariate]
+filter_funcs = [f_standard, f_standard_batched]  # , f_cholesky, f_univariate]
 
 filter_names = [
     "StandardFilter",
-    "CholeskyFilter",
-    "UnivariateFilter",
+    "StandardFilterBatched",
+    # "CholeskyFilter",
+    # "UnivariateFilter",
 ]
 
 output_names = [
@@ -65,17 +68,21 @@ def test_base_class_update_raises():
         filter.update(*inputs)
 
 
-@pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
-def test_output_shapes_one_state_one_observed(filter_func, rng):
+@pytest.mark.parametrize(
+    "filter_func, filter_name", zip(filter_funcs, filter_names), ids=filter_names
+)
+def test_output_shapes_one_state_one_observed(filter_func, filter_name, rng):
+    batch_size = 3 if "batched" in filter_name.lower() else 0
     p, m, r, n = 1, 1, 1, 10
-    inputs = make_test_inputs(p, m, r, n, rng)
-    outputs = filter_func(*inputs)
+    inputs = make_test_inputs(p, m, r, n, rng, batch_size=batch_size)
+    assert 0
+    # outputs = filter_func(*inputs)
 
     for output_idx, name in enumerate(output_names):
-        expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        expected_shape = get_expected_shape(name, p, m, r, n, batch_size)
+        # assert outputs[output_idx].shape == expected_shape, (
+        #     f"Shape of {name} does not match expected"
+        # )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
@@ -86,9 +93,9 @@ def test_output_shapes_when_all_states_are_stochastic(filter_func, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
@@ -99,9 +106,9 @@ def test_output_shapes_when_some_states_are_deterministic(filter_func, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.fixture
@@ -161,9 +168,9 @@ def test_output_shapes_with_time_varying_matrices(f_standard_nd, rng):
 
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
@@ -175,9 +182,9 @@ def test_output_with_deterministic_observation_equation(filter_func, rng):
 
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize(
@@ -190,9 +197,9 @@ def test_output_with_multiple_observed(filter_func, filter_name, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize(
@@ -206,9 +213,9 @@ def test_missing_data(filter_func, filter_name, p, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)