Handle latest PyMC API

Author: ricardoV94

conda-envs/environment-test.yml

@@ -10,6 +10,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.13.0  # CI was failing to resolve
+  - pymc>=5.16.0  # CI was failing to resolve
   - blackjax
   - scikit-learn

conda-envs/windows-environment-test.yml

@@ -10,6 +10,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.13.0  # CI was failing to resolve
+  - pymc>=5.16.0  # CI was failing to resolve
   - blackjax
   - scikit-learn

pymc_experimental/distributions/continuous.py

@@ -19,7 +19,7 @@
 The imports from pymc are not fully replicated here: add imports as necessary.
 """
 
-from typing import List, Tuple, Union
+from typing import Tuple, Union
 
 import numpy as np
 import pytensor.tensor as pt
@@ -37,8 +37,7 @@
 
 class GenExtremeRV(RandomVariable):
     name: str = "Generalized Extreme Value"
-    ndim_supp: int = 0
-    ndims_params: List[int] = [0, 0, 0]
+    signature = "(),(),()->()"
     dtype: str = "floatX"
     _print_name: Tuple[str, str] = ("Generalized Extreme Value", "\\operatorname{GEV}")
 
@@ -275,7 +274,7 @@ def chi_dist(nu: TensorVariable, size: TensorVariable) -> TensorVariable:
 
     def __new__(cls, name, nu, **kwargs):
         if "observed" not in kwargs:
-            kwargs.setdefault("transform", transforms.log)
+            kwargs.setdefault("default_transform", transforms.log)
         return CustomDist(name, nu, dist=cls.chi_dist, class_name="Chi", **kwargs)
 
     @classmethod

pymc_experimental/distributions/discrete.py

@@ -31,8 +31,7 @@ def log1mexp(x):
 
 class GeneralizedPoissonRV(RandomVariable):
     name = "generalized_poisson"
-    ndim_supp = 0
-    ndims_params = [0, 0]
+    signature = "(),()->()"
     dtype = "int64"
     _print_name = ("GeneralizedPoisson", "\\operatorname{GeneralizedPoisson}")
 

pymc_experimental/model/marginal_model.py

@@ -21,7 +21,7 @@
 from pytensor.graph.replace import graph_replace, vectorize_graph
 from pytensor.scan import map as scan_map
 from pytensor.tensor import TensorType, TensorVariable
-from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.shape import Shape
 from pytensor.tensor.special import log_softmax
 
@@ -598,7 +598,18 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs):
     fg = FunctionGraph(outputs=output_rvs, clone=False)
 
     non_elemwise_blockers = [
-        o for node in fg.apply_nodes if not isinstance(node.op, Elemwise) for o in node.outputs
+        o
+        for node in fg.apply_nodes
+        if not (
+            isinstance(node.op, Elemwise)
+            # Allow expand_dims on the left
+            or (
+                isinstance(node.op, DimShuffle)
+                and not node.op.drop
+                and node.op.shuffle == sorted(node.op.shuffle)
+            )
+        )
+        for o in node.outputs
     ]
     blocker_candidates = [rv_to_marginalize] + other_input_rvs + non_elemwise_blockers
     blockers = [var for var in blocker_candidates if var not in output_rvs]
@@ -698,16 +709,17 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
 
 def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> tuple[int, ...]:
     op = rv.owner.op
+    dist_params = rv.owner.op.dist_params(rv.owner)
     if isinstance(op, Bernoulli):
         return (0, 1)
     elif isinstance(op, Categorical):
-        p_param = rv.owner.inputs[3]
+        [p_param] = dist_params
         return tuple(range(pt.get_vector_length(p_param)))
     elif isinstance(op, DiscreteUniform):
-        lower, upper = constant_fold(rv.owner.inputs[3:])
+        lower, upper = constant_fold(dist_params)
         return tuple(np.arange(lower, upper + 1))
     elif isinstance(op, DiscreteMarkovChain):
-        P = rv.owner.inputs[0]
+        P, *_ = dist_params
         return tuple(range(pt.get_vector_length(P[-1])))
 
     raise NotImplementedError(f"Cannot compute domain for op {op}")

pymc_experimental/tests/model/test_marginal_model.py

@@ -7,10 +7,10 @@
 import pytensor.tensor as pt
 import pytest
 from arviz import InferenceData, dict_to_dataset
-from pymc import ImputationWarning, inputvars
 from pymc.distributions import transforms
 from pymc.logprob.abstract import _logprob
 from pymc.model.fgraph import fgraph_from_model
+from pymc.pytensorf import inputvars
 from pymc.util import UNSET
 from scipy.special import log_softmax, logsumexp
 from scipy.stats import halfnorm, norm
@@ -45,9 +45,7 @@ def disaster_model():
         early_rate = pm.Exponential("early_rate", 1.0, initval=3)
         late_rate = pm.Exponential("late_rate", 1.0, initval=1)
         rate = pm.math.switch(switchpoint >= years, early_rate, late_rate)
-        with pytest.warns(ImputationWarning), pytest.warns(
-            RuntimeWarning, match="invalid value encountered in cast"
-        ):
+        with pytest.warns(Warning):
             disasters = pm.Poisson("disasters", rate, observed=disaster_data)
 
     return disaster_model, years
@@ -294,7 +292,7 @@ def test_recover_marginals_basic():
 
     with m:
         prior = pm.sample_prior_predictive(
-            samples=20,
+            draws=20,
             random_seed=rng,
             return_inferencedata=False,
         )
@@ -337,7 +335,7 @@ def test_recover_marginals_coords():
 
     with m:
         prior = pm.sample_prior_predictive(
-            samples=20,
+            draws=20,
             random_seed=rng,
             return_inferencedata=False,
         )
@@ -364,7 +362,7 @@ def test_recover_batched_marginal():
 
     with m:
         prior = pm.sample_prior_predictive(
-            samples=20,
+            draws=20,
             random_seed=rng,
             return_inferencedata=False,
         )
@@ -394,7 +392,7 @@ def test_nested_recover_marginals():
 
     with m:
         prior = pm.sample_prior_predictive(
-            samples=20,
+            draws=20,
             random_seed=rng,
             return_inferencedata=False,
         )
@@ -565,7 +563,7 @@ def test_marginalized_transforms(transform, expected_warning):
             w=w,
             comp_dists=pm.HalfNormal.dist([1, 2, 3]),
             initval=initval,
-            transform=transform,
+            default_transform=transform,
         )
         y = pm.Normal("y", 0, sigma, observed=data)
 
@@ -583,7 +581,7 @@ def test_marginalized_transforms(transform, expected_warning):
                 ),
             ),
             initval=initval,
-            transform=transform,
+            default_transform=transform,
         )
         y = pm.Normal("y", 0, sigma, observed=data)
 

requirements.txt

@@ -1,2 +1,2 @@
-pymc>=5.13.0
+pymc>=5.16.0
 scikit-learn