Draft example notebook

Author: jessegrabowski

pymc_experimental/statespace/__init__.py

@@ -1,7 +1,7 @@
+from pymc_experimental.statespace.core.compile import compile_statespace
 from pymc_experimental.statespace.models import structural
 from pymc_experimental.statespace.models.ETS import BayesianETS
 from pymc_experimental.statespace.models.SARIMAX import BayesianSARIMA
 from pymc_experimental.statespace.models.VARMAX import BayesianVARMAX
-from pymc_experimental.statespace.utils import compile_statespace
 
 __all__ = ["structural", "BayesianSARIMA", "BayesianVARMAX", "BayesianETS", "compile_statespace"]

pymc_experimental/statespace/core/__init__.py

@@ -1,4 +1,5 @@
+from pymc_experimental.statespace.core.compile import compile_statespace
 from pymc_experimental.statespace.core.representation import PytensorRepresentation
 from pymc_experimental.statespace.core.statespace import PyMCStateSpace
 
-__all__ = ["PytensorRepresentation", "PyMCStateSpace"]
+__all__ = ["PytensorRepresentation", "PyMCStateSpace", "compile_statespace"]

pymc_experimental/statespace/utils/compile.py renamed to pymc_experimental/statespace/core/compile.py

@@ -461,7 +461,7 @@ def add_default_priors(self) -> None:
 
     def make_and_register_variable(
         self, name, shape: int | tuple[int, ...] | None = None, dtype=floatX
-    ) -> Variable:
+    ) -> pt.TensorVariable:
         """
         Helper function to create a pytensor symbolic variable and register it in the _name_to_variable dictionary
 

pymc_experimental/statespace/core/statespace.py

@@ -334,8 +334,8 @@ def param_dims(self):
             "seasonal_ma_params": (SEASONAL_MA_PARAM_DIM,),
         }
         if self.k_endog == 1:
-            coord_map["sigma_state"] = ()
-            coord_map["sigma_obs"] = ()
+            coord_map["sigma_state"] = None
+            coord_map["sigma_obs"] = None
         if not self.measurement_error:
             del coord_map["sigma_obs"]
         if self.p == 0:

pymc_experimental/statespace/models/SARIMAX.py

@@ -1,3 +0,0 @@
-from pymc_experimental.statespace.utils.compile import compile_statespace
-
-__all__ = ["compile_statespace"]