Yet more import fixes

Author: fonnesbeck

pymc_extras/distributions/multivariate/__init__.py

@@ -1,3 +1,3 @@
-from pymc_experimental.distributions.multivariate.r2d2m2cp import R2D2M2CP
+from pymc_extras.distributions.multivariate.r2d2m2cp import R2D2M2CP
 
 __all__ = ["R2D2M2CP"]

pymc_extras/inference/fit.py

@@ -34,11 +34,11 @@ def fit(method, **kwargs):
         if find_spec("blackjax") is None:
             raise RuntimeError("Need BlackJAX to use `pathfinder`")
 
-        from pymc_experimental.inference.pathfinder import fit_pathfinder
+        from pymc_extras.inference.pathfinder import fit_pathfinder
 
         return fit_pathfinder(**kwargs)
 
     if method == "laplace":
-        from pymc_experimental.inference.laplace import fit_laplace
+        from pymc_extras.inference.laplace import fit_laplace
 
         return fit_laplace(**kwargs)

pymc_extras/model/model_api.py

@@ -1,6 +1,7 @@
 from functools import wraps
 
 from pymc import Model
+import pymc_extras as pmx
 
 
 def as_model(*model_args, **model_kwargs):
@@ -16,7 +17,7 @@ def as_model(*model_args, **model_kwargs):
     .. code:: python
 
         import pymc as pm
-        import pymc_experimental as pmx
+        import pymc_extras as pmx
 
         # The following are equivalent
 

pymc_extras/statespace/core/__init__.py

@@ -1,7 +1,7 @@
 # ruff: noqa: I001
 
-from pymc_experimental.statespace.core.representation import PytensorRepresentation
-from pymc_experimental.statespace.core.statespace import PyMCStateSpace
-from pymc_experimental.statespace.core.compile import compile_statespace
+from pymc_extras.statespace.core.representation import PytensorRepresentation
+from pymc_extras.statespace.core.statespace import PyMCStateSpace
+from pymc_extras.statespace.core.compile import compile_statespace
 
 __all__ = ["PytensorRepresentation", "PyMCStateSpace", "compile_statespace"]

pymc_extras/statespace/filters/__init__.py

@@ -1,15 +1,19 @@
-from pymc_experimental.statespace.filters.distributions import LinearGaussianStateSpace
-from pymc_experimental.statespace.filters.kalman_filter import (
-    SquareRootFilter,
+from pymc_extras.statespace.filters.distributions import LinearGaussianStateSpace
+from pymc_extras.statespace.filters.kalman_filter import (
+    CholeskyFilter,
+    SingleTimeseriesFilter,
     StandardFilter,
+    SteadyStateFilter,
     UnivariateFilter,
 )
-from pymc_experimental.statespace.filters.kalman_smoother import KalmanSmoother
+from pymc_extras.statespace.filters.kalman_smoother import KalmanSmoother
 
 __all__ = [
     "StandardFilter",
     "UnivariateFilter",
+    "SteadyStateFilter",
     "KalmanSmoother",
-    "SquareRootFilter",
+    "SingleTimeseriesFilter",
+    "CholeskyFilter",
     "LinearGaussianStateSpace",
 ]

pymc_extras/statespace/models/__init__.py

@@ -1,6 +1,6 @@
-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_extras.statespace.models import structural
+from pymc_extras.statespace.models.ETS import BayesianETS
+from pymc_extras.statespace.models.SARIMAX import BayesianSARIMA
+from pymc_extras.statespace.models.VARMAX import BayesianVARMAX
 
 __all__ = ["structural", "BayesianSARIMA", "BayesianVARMAX", "BayesianETS"]