pymc-extras 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

pymc_extras/model/model_api.py

@@ -1,6 +1,9 @@
 from functools import wraps
+from inspect import signature
 
-from pymc import Model
+import pytensor.tensor as pt
+
+from pymc import Data, Model
 
 
 def as_model(*model_args, **model_kwargs):
@@ -9,6 +12,8 @@ def as_model(*model_args, **model_kwargs):
     This removes all need to think about context managers and lets you separate creating a generative model from using the model.
     Additionally, a coords argument is added to the function so coords can be changed during function invocation
 
+    All parameters are wrapped with a `pm.Data` object if the underlying type of the data supports it.
+
     Adapted from `Rob Zinkov's blog post <https://www.zinkov.com/posts/2023-alternative-frontends-pymc/>`_ and inspired by the `sampled <https://github.com/colcarroll/sampled>`_ decorator for PyMC3.
 
     Examples
@@ -47,8 +52,19 @@ def as_model(*model_args, **model_kwargs):
         @wraps(f)
         def make_model(*args, **kwargs):
             coords = model_kwargs.pop("coords", {}) | kwargs.pop("coords", {})
+            sig = signature(f)
+            ba = sig.bind(*args, **kwargs)
+            ba.apply_defaults()
+
             with Model(*model_args, coords=coords, **model_kwargs) as m:
-                f(*args, **kwargs)
+                for name, v in ba.arguments.items():
+                    # Only wrap pm.Data around values pytensor can process
+                    try:
+                        _ = pt.as_tensor_variable(v)
+                        ba.arguments[name] = Data(name, v)
+                    except (NotImplementedError, TypeError, ValueError):
+                        pass
+                f(*ba.args, **ba.kwargs)
             return m
 
         return make_model

pymc_extras/statespace/core/statespace.py

@@ -15,6 +15,9 @@ from pymc.model.transform.optimization import freeze_dims_and_data
 from pymc.util import RandomState
 from pytensor import Variable, graph_replace
 from pytensor.compile import get_mode
+from rich.box import SIMPLE_HEAD
+from rich.console import Console
+from rich.table import Table
 
 from pymc_extras.statespace.core.representation import PytensorRepresentation
 from pymc_extras.statespace.filters import (
@@ -254,53 +257,72 @@ class PyMCStateSpace:
         self.kalman_smoother = KalmanSmoother()
         self.make_symbolic_graph()
 
-        if verbose:
-            # These are split into separate try-except blocks, because it will be quite rare of models to implement
-            # _print_data_requirements, but we still want to print the prior requirements.
-            try:
-                self._print_prior_requirements()
-            except NotImplementedError:
-                pass
-            try:
-                self._print_data_requirements()
-            except NotImplementedError:
-                pass
-
-    def _print_prior_requirements(self) -> None:
-        """
-        Prints a short report to the terminal about the priors needed for the model, including their names,
+        self.requirement_table = None
+        self._populate_prior_requirements()
+        self._populate_data_requirements()
+
+        if verbose and self.requirement_table:
+            console = Console()
+            console.print(self.requirement_table)
+
+    def _populate_prior_requirements(self) -> None:
+        """
+        Add requirements about priors needed for the model to a rich table, including their names,
         shapes, named dimensions, and any parameter constraints.
         """
-        out = ""
-        for param, info in self.param_info.items():
-            out += f'\t{param} -- shape: {info["shape"]}, constraints: {info["constraints"]}, dims: {info["dims"]}\n'
-        out = out.rstrip()
+        # Check that the param_info class is implemented, and also that it's a dictionary. We can't proceed if either
+        # is not true.
+        try:
+            if not isinstance(self.param_info, dict):
+                return
+        except NotImplementedError:
+            return
 
-        _log.info(
-            "The following parameters should be assigned priors inside a PyMC "
-            f"model block: \n"
-            f"{out}"
-        )
+        if self.requirement_table is None:
+            self._initialize_requirement_table()
+
+        for param, info in self.param_info.items():
+            self.requirement_table.add_row(
+                param, str(info["shape"]), info["constraints"], str(info["dims"])
+            )
 
-    def _print_data_requirements(self) -> None:
+    def _populate_data_requirements(self) -> None:
         """
-        Prints a short report to the terminal about the data needed for the model, including their names, shapes,
-        and named dimensions.
+        Add requirements about the data needed for the model, including their names, shapes, and named dimensions.
         """
-        if not self.data_info:
+        try:
+            if not isinstance(self.data_info, dict):
+                return
+        except NotImplementedError:
             return
 
-        out = ""
+        if self.requirement_table is None:
+            self._initialize_requirement_table()
+        else:
+            self.requirement_table.add_section()
+
         for data, info in self.data_info.items():
-            out += f'\t{data} -- shape: {info["shape"]}, dims: {info["dims"]}\n'
-        out = out.rstrip()
+            self.requirement_table.add_row(data, str(info["shape"]), "pm.Data", str(info["dims"]))
+
+    def _initialize_requirement_table(self) -> None:
+        self.requirement_table = Table(
+            show_header=True,
+            show_edge=True,
+            box=SIMPLE_HEAD,
+            highlight=True,
+        )
 
-        _log.info(
-            "The following Data variables should be assigned to the model inside a PyMC "
-            f"model block: \n"
-            f"{out}"
+        self.requirement_table.title = "Model Requirements"
+        self.requirement_table.caption = (
+            "These parameters should be assigned priors inside a PyMC model block before "
+            "calling the build_statespace_graph method."
         )
 
+        self.requirement_table.add_column("Variable", justify="left")
+        self.requirement_table.add_column("Shape", justify="left")
+        self.requirement_table.add_column("Constraints", justify="left")
+        self.requirement_table.add_column("Dimensions", justify="right")
+
     def _unpack_statespace_with_placeholders(
         self,
     ) -> tuple[
@@ -961,10 +983,31 @@ class PyMCStateSpace:
         list[pm.Flat]
             A list of pm.Flat variables representing all parameters estimated by the model.
         """
+
+        def infer_variable_shape(name):
+            shape = self._name_to_variable[name].type.shape
+            if not any(dim is None for dim in shape):
+                return shape
+
+            dim_names = self._fit_dims.get(name, None)
+            if dim_names is None:
+                raise ValueError(
+                    f"Could not infer shape for {name}, because it was not given coords during model"
+                    f"fitting"
+                )
+
+            shape_from_coords = tuple([len(self._fit_coords[dim]) for dim in dim_names])
+            return tuple(
+                [
+                    shape[i] if shape[i] is not None else shape_from_coords[i]
+                    for i in range(len(shape))
+                ]
+            )
+
         for name in self.param_names:
             pm.Flat(
                 name,
-                shape=self._name_to_variable[name].type.shape,
+                shape=infer_variable_shape(name),
                 dims=self._fit_dims.get(name, None),
             )
 

pymc_extras/statespace/models/structural.py

@@ -1071,6 +1071,12 @@ class TimeSeasonality(Component):
 
         If None, states will be numbered ``[State_0, ..., State_s]``
 
+    remove_first_state: bool, default True
+        If True, the first state will be removed from the model. This is done because there are only n-1 degrees of
+        freedom in the seasonal component, and one state is not identified. If False, the first state will be
+        included in the model, but it will not be identified -- you will need to handle this in the priors (e.g. with
+        ZeroSumNormal).
+
     Notes
     -----
     A seasonal effect is any pattern that repeats every fixed interval. Although there are many possible ways to
@@ -1163,7 +1169,7 @@ class TimeSeasonality(Component):
         innovations: bool = True,
         name: str | None = None,
         state_names: list | None = None,
-        pop_state: bool = True,
+        remove_first_state: bool = True,
     ):
         if name is None:
             name = f"Seasonal[s={season_length}]"
@@ -1176,14 +1182,15 @@ class TimeSeasonality(Component):
                 )
             state_names = state_names.copy()
         self.innovations = innovations
-        self.pop_state = pop_state
+        self.remove_first_state = remove_first_state
 
-        if self.pop_state:
+        if self.remove_first_state:
             # In traditional models, the first state isn't identified, so we can help out the user by automatically
             # discarding it.
             # TODO: Can this be stashed and reconstructed automatically somehow?
             state_names.pop(0)
-            k_states = season_length - 1
+
+        k_states = season_length - int(self.remove_first_state)
 
         super().__init__(
             name=name,
@@ -1218,8 +1225,16 @@ class TimeSeasonality(Component):
             self.shock_names = [f"{self.name}"]
 
     def make_symbolic_graph(self) -> None:
-        T = np.eye(self.k_states, k=-1)
-        T[0, :] = -1
+        if self.remove_first_state:
+            # In this case, parameters are normalized to sum to zero, so the current state is the negative sum of
+            # all previous states.
+            T = np.eye(self.k_states, k=-1)
+            T[0, :] = -1
+        else:
+            # In this case we assume the user to be responsible for ensuring the states sum to zero, so T is just a
+            # circulant matrix that cycles between the states.
+            T = np.eye(self.k_states, k=1)
+            T[-1, 0] = 1
 
         self.ssm["transition", :, :] = T
         self.ssm["design", 0, 0] = 1

pymc_extras/utils/model_equivalence.py

@@ -0,0 +1,66 @@
+from collections.abc import Sequence
+
+from pymc.model.core import Model
+from pymc.model.fgraph import fgraph_from_model
+from pytensor import Variable
+from pytensor.compile import SharedVariable
+from pytensor.graph import Constant, graph_inputs
+from pytensor.graph.basic import equal_computations
+from pytensor.tensor.random.type import RandomType
+
+
+def equal_computations_up_to_root(
+    xs: Sequence[Variable], ys: Sequence[Variable], ignore_rng_values=True
+) -> bool:
+    # Check if graphs are equivalent even if root variables have distinct identities
+
+    x_graph_inputs = [var for var in graph_inputs(xs) if not isinstance(var, Constant)]
+    y_graph_inputs = [var for var in graph_inputs(ys) if not isinstance(var, Constant)]
+    if len(x_graph_inputs) != len(y_graph_inputs):
+        return False
+    for x, y in zip(x_graph_inputs, y_graph_inputs):
+        if x.type != y.type:
+            return False
+        if x.name != y.name:
+            return False
+        if isinstance(x, SharedVariable):
+            # if not isinstance(y, SharedVariable):
+            #     return False
+            if isinstance(x.type, RandomType) and ignore_rng_values:
+                continue
+            if not x.type.values_eq(x.get_value(), y.get_value()):
+                return False
+
+    return equal_computations(xs, ys, in_xs=x_graph_inputs, in_ys=y_graph_inputs)
+
+
+def equivalent_models(model1: Model, model2: Model) -> bool:
+    """Check whether two PyMC models are equivalent.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        import pymc as pm
+        from pymc_extras.utils.model_equivalence import equivalent_models
+
+        with pm.Model() as m1:
+            x = pm.Normal("x")
+            y = pm.Normal("y", x)
+
+        with pm.Model() as m2:
+            x = pm.Normal("x")
+            y = pm.Normal("y", x + 1)
+
+        with pm.Model() as m3:
+            x = pm.Normal("x")
+            y = pm.Normal("y", x)
+
+        assert not equivalent_models(m1, m2)
+        assert equivalent_models(m1, m3)
+
+    """
+    fgraph1, _ = fgraph_from_model(model1)
+    fgraph2, _ = fgraph_from_model(model2)
+    return equal_computations_up_to_root(fgraph1.outputs, fgraph2.outputs)

pymc_extras/version.txt

@@ -1 +1 @@
-0.2.0
+0.2.2

{pymc_extras-0.2.0.dist-info → pymc_extras-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: pymc-extras
-Version: 0.2.0
+Version: 0.2.2
 Summary: A home for new additions to PyMC, which may include unusual probability distribitions, advanced model fitting algorithms, or any code that may be inappropriate to include in the pymc repository, but may want to be made available to users.
 Home-page: http://github.com/pymc-devs/pymc-extras
 Maintainer: PyMC Developers
@@ -34,6 +34,17 @@ Provides-Extra: dev
 Requires-Dist: dask[all]; extra == "dev"
 Requires-Dist: blackjax; extra == "dev"
 Requires-Dist: statsmodels; extra == "dev"
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: maintainer
+Dynamic: maintainer-email
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # Welcome to `pymc-extras`
 <a href="https://gitpod.io/#https://github.com/pymc-devs/pymc-extras">
@@ -63,10 +74,9 @@ import pymc as pm
 import pymc_extras as pmx
 
 with pm.Model():
+  alpha = pmx.ParabolicFractal('alpha', b=1, c=1)
 
-    alpha = pmx.ParabolicFractal('alpha', b=1, c=1)
-
-    ...
+  ...
 
 ```
 

{pymc_extras-0.2.0.dist-info → pymc_extras-0.2.2.dist-info}/RECORD

@@ -1,31 +1,34 @@
-pymc_extras/__init__.py,sha256=WpDTZvLhxFg_t9gOE_wOSsswoYGIZpllsJH-_yOLEYI,1124
+pymc_extras/__init__.py,sha256=URh185f6b1xp2Taj2W2NJuW_hErKufBcLeQ0WDCyaNk,1160
 pymc_extras/linearmodel.py,sha256=6eitl15Ec15mSZu7zoHZ7Wwy4U1DPwqfAgwEt6ILeIc,3920
 pymc_extras/model_builder.py,sha256=sAw77fxdiy046BvDPjocuMlbJ0Efj-CDAGtmcwYmoG0,26361
 pymc_extras/printing.py,sha256=G8mj9dRd6i0PcsbcEWZm56ek6V8mmil78RI4MUhywBs,6506
 pymc_extras/version.py,sha256=VxPGCBzhtSegu-Jp5cjzn0n4DGU0wuPUh-KyZKB6uPM,240
-pymc_extras/version.txt,sha256=H5MN0fEzwfl6lP46y42zQ3LPTAH_2ys_9Mpy-UlBIek,6
+pymc_extras/version.txt,sha256=mY9riH7Xpu9E6EIQ0CN7cVvtQupyVPSNxBiv7ApIQQM,6
 pymc_extras/distributions/__init__.py,sha256=gTX7tvX8NcgP7V72URV7GeqF1aAEjGVbuW8LMxhXceY,1295
 pymc_extras/distributions/continuous.py,sha256=z-nvQgGncYISdRY8cWsa-56V0bQGq70jYwU-i8VZ0Uk,11253
 pymc_extras/distributions/discrete.py,sha256=vrARNuiQAEXrs7yQgImV1PO8AV1uyEC_LBhr6F9IcOg,13032
 pymc_extras/distributions/histogram_utils.py,sha256=5RTvlGCUrp2qzshrchmPyWxjhs6RIYL62SMikjDM1TU,5814
-pymc_extras/distributions/timeseries.py,sha256=EJxWOrfuQlODwPN13Udgy2ras6vQKS0Ebus0pUuduaA,12680
+pymc_extras/distributions/timeseries.py,sha256=M5MZ-nik_tgkaoZ1hdUGEZ9g04DQyVLwszVJqSKwNcY,12719
 pymc_extras/distributions/multivariate/__init__.py,sha256=E8OeLW9tTotCbrUjEo4um76-_WQD56PehsPzkKmhfyA,93
 pymc_extras/distributions/multivariate/r2d2m2cp.py,sha256=bUj9bB-hQi6CpaJfvJjgNPi727uTbvAdxl9fm1zNBqY,16005
 pymc_extras/gp/__init__.py,sha256=sFHw2y3lEl5tG_FDQHZUonQ_k0DF1JRf0Rp8dpHmge0,745
 pymc_extras/gp/latent_approx.py,sha256=cDEMM6H1BL2qyKg7BZU-ISrKn2HJe7hDaM4Y8GgQDf4,6682
 pymc_extras/inference/__init__.py,sha256=5cXpaQQnW0mJJ3x8wSxmYu63l--Xab5D_gMtjA6Q3uU,666
 pymc_extras/inference/find_map.py,sha256=T0uO8prUI5aBNuR1AN8fbA4cHmLRQLXznwJrfxfe7CA,15723
-pymc_extras/inference/fit.py,sha256=NFEpUaYLJAmDRP1WIPymgnEcXUofkoURYHbEdiTivzQ,1313
+pymc_extras/inference/fit.py,sha256=S9R48dh74s6K0MC9Iys4NAwVjP6rVRfx6SF-kPiR70E,1165
 pymc_extras/inference/laplace.py,sha256=OglOvnxfHLe0VXxBC1-ddVzADR9zgGxUPScM6P6FYo8,21163
-pymc_extras/inference/pathfinder.py,sha256=cmzR2OZCfkdTipT-8pmLuF-MHmLzxotsYlezOWBUM4U,4171
+pymc_extras/inference/pathfinder/__init__.py,sha256=FhAYrCWNx_dCrynEdjg2CZ9tIinvcVLBm67pNx_Y3kA,101
+pymc_extras/inference/pathfinder/importance_sampling.py,sha256=VvmuaE3aw_Mo3tMwswfF0rqe19mnhOCpzIScaJzjA1Y,6159
+pymc_extras/inference/pathfinder/lbfgs.py,sha256=P0UIOVtspdLzDU6alK-y91qzVAzXjYAXPuGmZ1nRqMo,5715
+pymc_extras/inference/pathfinder/pathfinder.py,sha256=fomZ5voVcWxvhWpeIZV7IHGIJCasT1g0ivC4dC3-0GM,63694
 pymc_extras/inference/smc/__init__.py,sha256=wyaT4NJl1YsSQRLiDy-i0Jq3CbJZ2BQd4nnCk-dIngY,603
 pymc_extras/inference/smc/sampling.py,sha256=AYwmKqGoV6pBtKnh9SUbBKbN7VcoFgb3MmNWV7SivMA,15365
 pymc_extras/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/model/model_api.py,sha256=_r6rYQG1tt9Z95QU-jVyHqZ1rs-u7sFMO5HJ5unDV5A,1750
+pymc_extras/model/model_api.py,sha256=UHMfQXxWBujeSiUySU0fDUC5Sd_BjT8FoVz3iBxQH_4,2400
 pymc_extras/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/model/marginal/distributions.py,sha256=hytO-CJqMXuvIfJhdfKNlf9tkblOwCp8hAWTiIS0cyU,12088
-pymc_extras/model/marginal/graph_analysis.py,sha256=LNx5N8ZE7d8Sq3BnlUYHrrwnxTLJTvehf8xYu95yrb8,15699
-pymc_extras/model/marginal/marginal_model.py,sha256=5DbhjlOAwY6JSMJUhUAPXPpybXQU0x7MzOGd3eXACYo,24854
+pymc_extras/model/marginal/distributions.py,sha256=iM1yT7_BmivgUSloQPKE2QXGPgjvLqDMY_OTBGsdAWg,15563
+pymc_extras/model/marginal/graph_analysis.py,sha256=0hWUH_PjfpgneQ3NaT__pWHS1fh50zNbI86kH4Nub0E,15693
+pymc_extras/model/marginal/marginal_model.py,sha256=oNsiSWHjOPCTDxNEivEILLP_cOuBarm29Gr2p6hWHIM,23594
 pymc_extras/model/transforms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/transforms/autoreparam.py,sha256=_NltGWmNqi_X9sHCqAvWcBveLTPxVy11-wENFTcN6kk,12377
 pymc_extras/preprocessing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -34,7 +37,7 @@ pymc_extras/statespace/__init__.py,sha256=0MtZj7yT6jcyERvITnn-nkhyY8fO6Za4_vV53C
 pymc_extras/statespace/core/__init__.py,sha256=huHEiXAm8zV2MZyZ8GBHp6q7_fnWqveM7lC6ilpb3iE,309
 pymc_extras/statespace/core/compile.py,sha256=1c8Q9D9zeUe7F0z7CH6q1C6ZuLg2_imgk8RoE_KMaFI,1608
 pymc_extras/statespace/core/representation.py,sha256=DwNIun6wdeEA20oWBx5M4govyWTf5JI87aGQ_E6Mb4U,18956
-pymc_extras/statespace/core/statespace.py,sha256=ZElRm9wJvIGG4Pw-3qiQpBkHXRDqS6pfRyuGrBBcZ2Y,95270
+pymc_extras/statespace/core/statespace.py,sha256=K_WVnWKlI6sR2kgriq9sctQVvwXCeAirm14TthDpmRM,96860
 pymc_extras/statespace/filters/__init__.py,sha256=N9Q4D0gAq_ZtT-GtrqiX1HkSg6Orv7o1TbrWUtnbTJE,420
 pymc_extras/statespace/filters/distributions.py,sha256=-9j__vRqL5hKyYFnQr5HKHA5kEFzwiuSccH4mslTOuQ,12900
 pymc_extras/statespace/filters/kalman_filter.py,sha256=HELC3aK4k8EdWlUAk5_F7y7YkIz-Xi_0j2AwRgAXgcc,31949
@@ -44,7 +47,7 @@ pymc_extras/statespace/models/ETS.py,sha256=o039M-6aCxyMXbbKvUeNVZhheCKvvNIAmuj0
 pymc_extras/statespace/models/SARIMAX.py,sha256=SX0eiSK1pOt4dHBjWzBqVpRz67pBGLN5pQQgXcOiOgY,21607
 pymc_extras/statespace/models/VARMAX.py,sha256=xkIuftNc_5NHFpqZalExni99-1kovnzm5OjMIDNgaxY,15989
 pymc_extras/statespace/models/__init__.py,sha256=U79b8rTHBNijVvvGOd43nLu4PCloPUH1rwlN87-n88c,317
-pymc_extras/statespace/models/structural.py,sha256=W5FmImZyvHGxpFCYczfi6IVIjXDQkzaFjVKZSa6CiW8,63017
+pymc_extras/statespace/models/structural.py,sha256=sep9pesJdRN4X8Bea6_RhO3112uWOZRuYRxO6ibl_OA,63943
 pymc_extras/statespace/models/utilities.py,sha256=G9GuHKsghmIYOlfkPtvxBWF-FZY5-5JI1fJQM8N7EnE,15373
 pymc_extras/statespace/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/statespace/utils/constants.py,sha256=Kf6j75ABaDQeRODxKQ76wTUQV4F5sTjn1KBcZgCQx20,2403
@@ -52,6 +55,7 @@ pymc_extras/statespace/utils/coord_tools.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm
 pymc_extras/statespace/utils/data_tools.py,sha256=caanvrxDu9g-dEKff2bbmaVTs6-71kkSoYIiiSUXhw4,5985
 pymc_extras/utils/__init__.py,sha256=yxI9cJ7fCtVQS0GFw0y6mDGZIQZiK53vm3UNKqIuGSk,758
 pymc_extras/utils/linear_cg.py,sha256=KkXhuimFsrKtNd_0By2ApxQQQNm5FdBtmDQJOVbLYkA,10056
+pymc_extras/utils/model_equivalence.py,sha256=8QIftID2HDxD659i0RXHazQ-l2Q5YegCRLcDqb2p9Pc,2187
 pymc_extras/utils/pivoted_cholesky.py,sha256=QtnjP0pAl9b77fLAu-semwT4_9dcoiqx3dz1xKGBjMk,1871
 pymc_extras/utils/prior.py,sha256=QlWVr7uKIK9VncBw7Fz3YgaASKGDfqpORZHc-vz_9gQ,6841
 pymc_extras/utils/spline.py,sha256=qGq0gcoMG5dpdazKFzG0RXkkCWP8ADPPXN-653-oFn4,4820
@@ -62,23 +66,23 @@ tests/test_histogram_approximation.py,sha256=w-xb2Rr0Qft6sm6F3BTmXXnpuqyefC1SUL6
 tests/test_laplace.py,sha256=5ioEyP6AzmMszrtQRz0KWTsCCU35SEhSOdBcYfYzptE,8228
 tests/test_linearmodel.py,sha256=iB8ApNqIX9_nUHoo-Tm51xuPdrva5t4VLLut6qXB5Ao,6906
 tests/test_model_builder.py,sha256=QiINEihBR9rx8xM4Nqlg4urZKoyo58aTKDtxl9SJF1s,11249
-tests/test_pathfinder.py,sha256=FBm0ge6rje5jz9_10h_247E70aKCpkbu1jmzrR7Ar8A,1726
-tests/test_pivoted_cholesky.py,sha256=7_thrb90_an_S3boYr0mu4NNOhjiI6AaZ2ADn53sBX8,698
+tests/test_pathfinder.py,sha256=GnSbZJ9QuFW9UVbkWaVgMVqQZTCttOyz_rSflxhQ-EA,4955
+tests/test_pivoted_cholesky.py,sha256=PuMdMSCzO4KdQWpUF4SEBeuH_qsINCIH8TYtmmJ1NKo,692
 tests/test_printing.py,sha256=HnvwwjrjBuxXFAJdyU0K_lvKGLgh4nzHAnhsIUpenbY,5211
 tests/test_prior_from_trace.py,sha256=HOzR3l98pl7TEJquo_kSugED4wBTgHo4-8lgnpmacs8,5516
 tests/test_splines.py,sha256=xSZi4hqqReN1H8LHr0xjDmpomhDQm8auIsWQjFOyjbM,2608
-tests/utils.py,sha256=cRPe0ovsexxOQ6dK94xao_Kv5qcPrqtFWOFXBqubHqY,1257
+tests/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/distributions/__init__.py,sha256=jt-oloszTLNFwi9AgU3M4m6xKQ8xpQE338rmmaMZcMs,795
 tests/distributions/test_continuous.py,sha256=1-bu-IP6RgLUJnuPYpOD8ZS1ahYbKtsJ9oflBfqCaFo,5477
 tests/distributions/test_discrete.py,sha256=CjjaUpppsvQ6zLzV15ZsbwNOKrDmEdz4VWcleoCXUi0,7776
 tests/distributions/test_discrete_markov_chain.py,sha256=8RCHZXSB8IWjniuKaGGlM_iTWGmdrcOqginxmrAeEJg,9212
 tests/distributions/test_multivariate.py,sha256=LBvBuoT_3rzi8rR38b8L441Y-9Ff0cIXeRBKiEn6kjs,10452
 tests/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/model/test_model_api.py,sha256=SiOMA1NpyQKJ7stYI1ms8ksDPU81lVo8wS8hbqiik-U,776
+tests/model/test_model_api.py,sha256=FJvMTmexovRELZOUcUyk-6Vwk9qSiH7hIFoiArgl5mk,1040
 tests/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/model/marginal/test_distributions.py,sha256=UWoXi1h-e0QG0fLxPXkUj2LXzNb5uhoHicIBIVDIhL4,5126
+tests/model/marginal/test_distributions.py,sha256=p5f73g4ogxYkdZaBndZV_1ra8TCppXiRlUpaaTwEe-M,5195
 tests/model/marginal/test_graph_analysis.py,sha256=raoj41NusMOj1zzPCrxrlQODqX6Ey8Ft_o32pNTe5qg,6712
-tests/model/marginal/test_marginal_model.py,sha256=IWmF_XkbQeI15H0uepi1zET1Zffj6qVgAAdAiUoxrhA,31009
+tests/model/marginal/test_marginal_model.py,sha256=uOmARalkdWq3sDbnJQ0KjiLwviqauZOAnafmYS_Cnd8,35475
 tests/statespace/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/statespace/test_ETS.py,sha256=IPg3uQ7xEGqDMEHu993vtUTV7r-uNAxmw23sr5MVGfQ,15582
 tests/statespace/test_SARIMAX.py,sha256=1BYNOm9aSHnpn-qbpe3YsQVH8m-mXcp_gvKgWhWn1W4,12948
@@ -87,15 +91,15 @@ tests/statespace/test_coord_assignment.py,sha256=2GBm46-0eI4QNh4bvp3D7az58stcA5Z
 tests/statespace/test_distributions.py,sha256=WQ_ROyd-PL3cimXTyEtyVaMEVtS7Hue2Z0lN7UnGDyo,9122
 tests/statespace/test_kalman_filter.py,sha256=s2n62FzXl9elU_uqaMNaEaexUfq3SXe3_YvQ2lM6hiQ,11600
 tests/statespace/test_representation.py,sha256=1KAJY4ZaVhb1WdAJLx2UYSXuVYsMNWX98gEDF7P0B4s,6210
-tests/statespace/test_statespace.py,sha256=8ZLLQaxlP5UEJnIMYyIzzAODCxMxs6E5I1hLu2HCdqo,28866
+tests/statespace/test_statespace.py,sha256=JoupFFpG8PmpB_NFV471IuTmyXhEd6_vOISwVCRrBBM,30570
 tests/statespace/test_statespace_JAX.py,sha256=hZOc6xxYdVeATPCKmcHMLOVcuvdzGRzgQQ4RrDenwk8,5279
-tests/statespace/test_structural.py,sha256=IN6OQuq7bZVCDiws3Yhsa4IoPyfLaONzgzdvImp0Zcc,29036
+tests/statespace/test_structural.py,sha256=HD8OaGbjuH4y3xv_uG-R1xLZpPpcb4-3dbcTeb_imLY,29306
 tests/statespace/utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/statespace/utilities/shared_fixtures.py,sha256=SNw8Bvj1Yw11TxAW6n20Bq0B8oaYtVTiFFEVNH_wnp4,164
 tests/statespace/utilities/statsmodel_local_level.py,sha256=SQAzaYaSDwiVhUQ1iWjt4MgfAd54RuzVtnslIs3xdS8,1225
 tests/statespace/utilities/test_helpers.py,sha256=oH24a6Q45NFFFI3Kx9mhKbxsCvo9ErCorKFoTjDB3-4,9159
-pymc_extras-0.2.0.dist-info/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
-pymc_extras-0.2.0.dist-info/METADATA,sha256=DpLFVnYDJouXKp6hsIfTtE6rBgJfcLvRdq5CsEKaNVQ,4899
-pymc_extras-0.2.0.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-pymc_extras-0.2.0.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
-pymc_extras-0.2.0.dist-info/RECORD,,
+pymc_extras-0.2.2.dist-info/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
+pymc_extras-0.2.2.dist-info/METADATA,sha256=9k60kKKNzr7E24gACpTBNOWj-tRTNOiujSZOFD89G5c,5140
+pymc_extras-0.2.2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+pymc_extras-0.2.2.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
+pymc_extras-0.2.2.dist-info/RECORD,,

{pymc_extras-0.2.0.dist-info → pymc_extras-0.2.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tests/model/marginal/test_distributions.py

@@ -6,7 +6,7 @@ from pymc.logprob.abstract import _logprob
 from pytensor import tensor as pt
 from scipy.stats import norm
 
-from pymc_extras import MarginalModel
+from pymc_extras import marginalize
 from pymc_extras.distributions import DiscreteMarkovChain
 from pymc_extras.model.marginal.distributions import MarginalFiniteDiscreteRV
 
@@ -21,6 +21,7 @@ def test_marginalized_bernoulli_logp():
         [mu],
         [idx, y],
         dims_connections=(((),),),
+        dims=(),
     )(mu)[0].owner
 
     y_vv = y.clone()
@@ -43,7 +44,7 @@ def test_marginalized_hmm_normal_emission(batch_chain, batch_emission):
     if batch_chain and not batch_emission:
         pytest.skip("Redundant implicit combination")
 
-    with MarginalModel() as m:
+    with pm.Model() as m:
         P = [[0, 1], [1, 0]]
         init_dist = pm.Categorical.dist(p=[1, 0])
         chain = DiscreteMarkovChain(
@@ -53,8 +54,8 @@ def test_marginalized_hmm_normal_emission(batch_chain, batch_emission):
             "emission", mu=chain * 2 - 1, sigma=1e-1, shape=(3, 4) if batch_emission else None
         )
 
-    m.marginalize([chain])
-    logp_fn = m.compile_logp()
+    marginal_m = marginalize(m, [chain])
+    logp_fn = marginal_m.compile_logp()
 
     test_value = np.array([-1, 1, -1, 1])
     expected_logp = pm.logp(pm.Normal.dist(0, 1e-1), np.zeros_like(test_value)).sum().eval()
@@ -70,7 +71,7 @@ def test_marginalized_hmm_normal_emission(batch_chain, batch_emission):
 )
 def test_marginalized_hmm_categorical_emission(categorical_emission):
     """Example adapted from https://www.youtube.com/watch?v=9-sPm4CfcD0"""
-    with MarginalModel() as m:
+    with pm.Model() as m:
         P = np.array([[0.5, 0.5], [0.3, 0.7]])
         init_dist = pm.Categorical.dist(p=[0.375, 0.625])
         chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, steps=2)
@@ -78,11 +79,11 @@ def test_marginalized_hmm_categorical_emission(categorical_emission):
             emission = pm.Categorical("emission", p=pt.constant([[0.8, 0.2], [0.4, 0.6]])[chain])
         else:
             emission = pm.Bernoulli("emission", p=pt.where(pt.eq(chain, 0), 0.2, 0.6))
-    m.marginalize([chain])
+    marginal_m = marginalize(m, [chain])
 
     test_value = np.array([0, 0, 1])
     expected_logp = np.log(0.1344)  # Shown at the 10m22s mark in the video
-    logp_fn = m.compile_logp()
+    logp_fn = marginal_m.compile_logp()
     np.testing.assert_allclose(logp_fn({"emission": test_value}), expected_logp)
 
 
@@ -95,7 +96,7 @@ def test_marginalized_hmm_multiple_emissions(batch_chain, batch_emission1, batch
         (2, *reversed(chain_shape)) if batch_emission1 else tuple(reversed(chain_shape))
     )
     emission2_shape = (*chain_shape, 2) if batch_emission2 else chain_shape
-    with MarginalModel() as m:
+    with pm.Model() as m:
         P = [[0, 1], [1, 0]]
         init_dist = pm.Categorical.dist(p=[1, 0])
         chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, shape=chain_shape)
@@ -108,10 +109,10 @@ def test_marginalized_hmm_multiple_emissions(batch_chain, batch_emission1, batch
             emission2_mu = emission2_mu[..., None]
         emission_2 = pm.Normal("emission_2", mu=emission2_mu, sigma=1e-1, shape=emission2_shape)
 
-    with pytest.warns(UserWarning, match="multiple dependent variables"):
-        m.marginalize([chain])
+    marginal_m = marginalize(m, [chain])
 
-    logp_fn = m.compile_logp(sum=False)
+    with pytest.warns(UserWarning, match="multiple dependent variables"):
+        logp_fn = marginal_m.compile_logp(sum=False)
 
     test_value = np.array([-1, 1, -1, 1])
     multiplier = 2 + batch_emission1 + batch_emission2