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

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.1

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pymc-extras
-Version: 0.2.0
+Version: 0.2.1
 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
@@ -63,10 +63,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.1.dist-info}/RECORD

@@ -1,14 +1,14 @@
-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=cQFcl5zLD8igvnygroMEarBFzcLI-qCfsvD35ED5tKY,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
@@ -23,9 +23,9 @@ pymc_extras/inference/smc/sampling.py,sha256=AYwmKqGoV6pBtKnh9SUbBKbN7VcoFgb3MmN
 pymc_extras/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/model/model_api.py,sha256=_r6rYQG1tt9Z95QU-jVyHqZ1rs-u7sFMO5HJ5unDV5A,1750
 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
@@ -44,7 +44,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 +52,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
@@ -63,11 +64,11 @@ 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_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
@@ -76,9 +77,9 @@ tests/distributions/test_multivariate.py,sha256=LBvBuoT_3rzi8rR38b8L441Y-9Ff0cIX
 tests/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/model/test_model_api.py,sha256=SiOMA1NpyQKJ7stYI1ms8ksDPU81lVo8wS8hbqiik-U,776
 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
@@ -89,13 +90,13 @@ tests/statespace/test_kalman_filter.py,sha256=s2n62FzXl9elU_uqaMNaEaexUfq3SXe3_Y
 tests/statespace/test_representation.py,sha256=1KAJY4ZaVhb1WdAJLx2UYSXuVYsMNWX98gEDF7P0B4s,6210
 tests/statespace/test_statespace.py,sha256=8ZLLQaxlP5UEJnIMYyIzzAODCxMxs6E5I1hLu2HCdqo,28866
 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.1.dist-info/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
+pymc_extras-0.2.1.dist-info/METADATA,sha256=pT1MOjFxsX6lc0q_D3J-2jNW6UaRkCq_0kJemgG4DGU,4894
+pymc_extras-0.2.1.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
+pymc_extras-0.2.1.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
+pymc_extras-0.2.1.dist-info/RECORD,,

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

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.7.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