Vectorize statespace matrix builder

Author: aandorra-mia

pymc_extras/statespace/core/statespace.py

@@ -46,6 +46,7 @@
     VECTOR_VALUED,
 )
 from pymc_extras.statespace.utils.data_tools import register_data_with_pymc
+from pytensor.graph.replace import vectorize_graph
 
 _log = logging.getLogger("pymc.experimental.statespace")
 
@@ -726,7 +727,7 @@ def _insert_random_variables(self):
         matrices = list(self._unpack_statespace_with_placeholders())
 
         replacement_dict = {var: pymc_model[name] for name, var in self._name_to_variable.items()}
-        self.subbed_ssm = graph_replace(matrices, replace=replacement_dict, strict=True)
+        self.subbed_ssm = vectorize_graph(matrices, replace=replacement_dict)
 
     def _insert_data_variables(self):
         """

tests/statespace/test_statespace.py

@@ -710,8 +710,7 @@ def test_invalid_scenarios():
     # Giving a list, tuple, or Series when a matrix of data is expected should always raise
     with pytest.raises(
         ValueError,
-        match="Scenario data for variable 'a' has the wrong number of columns. "
-        "Expected 2, got 1",
+        match="Scenario data for variable 'a' has the wrong number of columns. Expected 2, got 1",
     ):
         for data_type in [list, tuple, pd.Series]:
             ss_mod._validate_scenario_data(data_type(np.zeros(10)))
@@ -720,15 +719,14 @@ def test_invalid_scenarios():
     # Providing irrevelant data raises
     with pytest.raises(
         ValueError,
-        match="Scenario data provided for variable 'jk lol', which is not an exogenous " "variable",
+        match="Scenario data provided for variable 'jk lol', which is not an exogenous variable",
     ):
         ss_mod._validate_scenario_data({"jk lol": np.zeros(10)})
 
     # Incorrect 2nd dimension of a non-dataframe
     with pytest.raises(
         ValueError,
-        match="Scenario data for variable 'a' has the wrong number of columns. Expected "
-        "2, got 1",
+        match="Scenario data for variable 'a' has the wrong number of columns. Expected 2, got 1",
     ):
         scenario = np.zeros(10).tolist()
         ss_mod._validate_scenario_data(scenario)
@@ -870,3 +868,13 @@ def test_forecast_with_exog_data(rng, exog_ss_mod, idata_exog, start):
     regression_effect_expected = (betas * scenario_xr).sum(dim=["state"])
 
     assert_allclose(regression_effect, regression_effect_expected)
+
+
+@pytest.mark.parametrize("batch_size", [(10,), (10, 3, 5)])
+def test_insert_batched_rvs(ss_mod, batch_size):
+    with pm.Model():
+        rho = pm.Normal("rho", shape=batch_size)
+        zeta = pm.Normal("zeta", shape=batch_size)
+        ss_mod._insert_random_variables()
+    matrices = ss_mod.unpack_statespace()
+    assert matrices[4].type.shape == (*batch_size, 2, 2)