start vectorizing inside scan

Author: aandorra-mia

pymc_extras/statespace/filters/kalman_filter.py

@@ -282,9 +282,10 @@ def _build_graph(
 
         if len(sequences) > 0:
             sequences = self.add_check_on_time_varying_shapes(data, sequences)
-
+# y, a, P, c, d, T, Z, R, H, Q = self.unpack_args(args)
+        # s0 and s1 = states, o0 and 01 are observed states, r0 and r1 are random states
         results, updates = pytensor.scan(
-            self.kalman_step,
+            pt.vectorize(self.kalman_step, signature="(o0,), (s0,), (s0, s1), (s0,), (o0,), (s0, s1), (o0, s0), (s0, r0), (o0, o1), (r0, r1) -> (s0,), (s0,), (o0,), (s0, s1), (P_hat,), (obs_cov,), (ll,)",
             sequences=[data, *sequences],
             outputs_info=[None, a0, None, None, P0, None, None],
             non_sequences=non_sequences,
@@ -320,18 +321,17 @@ def build_graph(
         """
         Build the vectorized computation graph for the Kalman filter.
         """
-        signature = self._make_gufunc_signature(
-            [data, a0, P0, c, d, T, Z, R, H, Q],
-        )
+        # signature = self._make_gufunc_signature(
+        #     [data, a0, P0, c, d, T, Z, R, H, Q],
+        # )
         fn = partial(
             self._build_graph,
             mode=mode,
             return_updates=return_updates,
             missing_fill_value=missing_fill_value,
             cov_jitter=cov_jitter,
         )
-        filter_outputs = pt.vectorize(fn, signature=signature)(data, a0, P0, c, d, T, Z, R, H, Q)
-        # filter_outputs = fn(data, a0, P0, c, d, T, Z, R, H, Q)
+        filter_outputs = fn(data, a0, P0, c, d, T, Z, R, H, Q)
         for output, name in zip(filter_outputs, ALL_KF_OUTPUT_NAMES):
             output.name = name
 

pymc_extras/statespace/filters/kalman_smoother.py

@@ -169,16 +169,15 @@ def build_graph(
         """
         Build the vectorized computation graph for the Kalman smoother.
         """
-        signature = self._make_gufunc_signature(
-            [T, R, Q, filtered_states, filtered_covariances],
-        )
+        # signature = self._make_gufunc_signature(
+        #     [T, R, Q, filtered_states, filtered_covariances],
+        # )
         fn = partial(
             self._build_graph,
             mode=mode,
             cov_jitter=cov_jitter,
         )
-        return pt.vectorize(fn, signature=signature)(T, R, Q, filtered_states, filtered_covariances)
-        # return fn(T, R, Q, filtered_states, filtered_covariances)
+        return fn(T, R, Q, filtered_states, filtered_covariances)
 
     def smoother_step(self, *args):
         a, P, a_smooth, P_smooth, T, R, Q = self.unpack_args(args)

tests/statespace/test_kalman_filter.py

@@ -336,6 +336,7 @@ def test_batched_standard_filter(filter_func):
         pt.as_tensor(x, name=name)
         for x, name in zip(make_test_inputs(p, m, r, n, rng, batch_size=8), input_names)
     ]
+    inputs[0] = pt.moveaxis(inputs[0], -2, 0)  # Move time dimension to the front of the data array
     kf = StandardFilter()
     outputs = kf.build_graph(*inputs)
     fn = pytensor.function([], outputs)