Update for PyMC v5.11 (#320)

Author: zaxtax

conda-envs/environment-test.yml

@@ -11,6 +11,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.10.0  # CI was failing to resolve
+  - pymc>=5.11.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.10.0  # CI was failing to resolve
+  - pymc>=5.11.0  # CI was failing to resolve
   - blackjax
   - scikit-learn

pymc_experimental/distributions/continuous.py

@@ -211,7 +211,7 @@ def logcdf(value, mu, sigma, xi):
             logc, sigma > 0, pt.and_(xi > -1, xi < 1), msg="sigma > 0 or -1 < xi < 1"
         )
 
-    def moment(rv, size, mu, sigma, xi):
+    def support_point(rv, size, mu, sigma, xi):
         r"""
         Using the mode, as the mean can be infinite when :math:`\xi > 1`
         """

pymc_experimental/distributions/discrete.py

@@ -143,7 +143,7 @@ def dist(cls, mu, lam, **kwargs):
         lam = pt.as_tensor_variable(lam)
         return super().dist([mu, lam], **kwargs)
 
-    def moment(rv, size, mu, lam):
+    def support_point(rv, size, mu, lam):
         mean = pt.floor(mu / (1 - lam))
         if not rv_size_is_none(size):
             mean = pt.full(size, mean)

pymc_experimental/distributions/timeseries.py

@@ -9,8 +9,8 @@
 from pymc.distributions.distribution import (
     Distribution,
     SymbolicRandomVariable,
-    _moment,
-    moment,
+    _support_point,
+    support_point,
 )
 from pymc.distributions.shape_utils import (
     _change_dist_size,
@@ -221,9 +221,9 @@ def change_mc_size(op, dist, new_size, expand=False):
     return DiscreteMarkovChain.rv_op(*dist.owner.inputs[:-1], size=new_size, n_lags=op.n_lags)
 
 
-@_moment.register(DiscreteMarkovChainRV)
+@_support_point.register(DiscreteMarkovChainRV)
 def discrete_mc_moment(op, rv, P, steps, init_dist, state_rng):
-    init_dist_moment = moment(init_dist)
+    init_dist_moment = support_point(init_dist)
     n_lags = op.n_lags
 
     def greedy_transition(*args):

pymc_experimental/model/marginal_model.py

@@ -655,7 +655,7 @@ def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> Tuple[int, ...]:
         return tuple(range(pt.get_vector_length(p_param)))
     elif isinstance(op, DiscreteUniform):
         lower, upper = constant_fold(rv.owner.inputs[3:])
-        return tuple(range(lower, upper + 1))
+        return tuple(np.arange(lower, upper + 1))
     elif isinstance(op, DiscreteMarkovChain):
         P = rv.owner.inputs[0]
         return tuple(range(pt.get_vector_length(P[-1])))

pymc_experimental/tests/distributions/test_continuous.py

@@ -25,7 +25,7 @@
     R,
     Rplus,
     Rplusbig,
-    assert_moment_is_expected,
+    assert_support_point_is_expected,
     check_logcdf,
     check_logp,
     seeded_scipy_distribution_builder,
@@ -111,10 +111,10 @@ def ref_logcdf(value, mu, sigma, xi):
             ),
         ],
     )
-    def test_genextreme_moment(self, mu, sigma, xi, size, expected):
+    def test_genextreme_support_point(self, mu, sigma, xi, size, expected):
         with pm.Model() as model:
             GenExtreme("x", mu=mu, sigma=sigma, xi=xi, size=size)
-        assert_moment_is_expected(model, expected)
+        assert_support_point_is_expected(model, expected)
 
     def test_gen_extreme_scipy_kwarg(self):
         dist = GenExtreme.dist(xi=1, scipy=False)

pymc_experimental/tests/distributions/test_discrete.py

@@ -23,7 +23,7 @@
     Domain,
     I,
     Rplus,
-    assert_moment_is_expected,
+    assert_support_point_is_expected,
     check_logp,
     discrete_random_tester,
 )
@@ -123,7 +123,7 @@ def test_logp_lam_expected_moments(self):
     def test_moment(self, mu, lam, size, expected):
         with pm.Model() as model:
             GeneralizedPoisson("x", mu=mu, lam=lam, size=size)
-        assert_moment_is_expected(model, expected)
+        assert_support_point_is_expected(model, expected)
 
 
 class TestBetaNegativeBinomial:

pymc_experimental/tests/distributions/test_discrete_markov_chain.py

@@ -129,7 +129,7 @@ def test_moment_function(self):
             state = chain_np[i]
             chain_np[i + 1] = np.argmax(P_np[state])
 
-        dmc_chain = pm.distributions.distribution.moment(chain).eval()
+        dmc_chain = pm.distributions.distribution.support_point(chain).eval()
 
         assert np.allclose(dmc_chain, chain_np)
 

pymc_experimental/tests/model/test_marginal_model.py

@@ -58,7 +58,7 @@ def test_marginalized_bernoulli_logp():
     marginal_rv_node = FiniteDiscreteMarginalRV(
         [mu],
         [idx, y],
-        ndim_supp=None,
+        ndim_supp=0,
         n_updates=0,
     )(
         mu