Fix call of logpt in test_dirichlet_multinomial_matches_beta_binomial

Author: AlexAndorra

pymc3/tests/test_distributions.py

@@ -2273,11 +2273,17 @@ def test_dirichlet_multinomial_matches_beta_binomial(self):
         a, b, n = 2, 1, 5
         ns = np.arange(n + 1)
         ns_dm = np.vstack((ns, n - ns)).T  # convert ns=1 to ns_dm=[1, 4], for all ns...
-        bb_logp = logpt(var=pm.BetaBinomial.dist(n=n, alpha=a, beta=b, size=2), rv_values=ns).eval()
-        dm_logp = logpt(
-            var=pm.DirichletMultinomial.dist(n=n, a=[a, b], size=2),
-            rv_values=ns_dm,
-        ).eval().ravel()
+
+        bb = pm.BetaBinomial.dist(n=n, alpha=a, beta=b, size=2)
+        bb_value = bb.type()
+        bb.tag.value_var = bb_value
+        bb_logp = logpt(var=bb, rv_values={bb: bb_value}).eval({bb_value: ns})
+
+        dm = pm.DirichletMultinomial.dist(n=n, a=[a, b], size=2)
+        dm_value = dm.type()
+        dm.tag.value_var = dm_value
+        dm_logp = logpt(var=dm, rv_values={dm: dm_value}).eval({dm_value: ns_dm}).ravel()
+
         assert_almost_equal(
             dm_logp,
             bb_logp,
@@ -2290,19 +2296,19 @@ def test_dirichlet_multinomial_vec(self):
         n = 10
 
         with Model() as model_single:
-            DirichletMultinomial("m", n=n, a=a)
+            pm.DirichletMultinomial("m", n=n, a=a)
 
         with Model() as model_many:
-            DirichletMultinomial("m", n=n, a=a, size=2)
+            pm.DirichletMultinomial("m", n=n, a=a, size=2)
 
         assert_almost_equal(
-            np.asarray([dirichlet_multinomial_logpmf(v, n, a) for v in vals]),
+            np.asarray([dirichlet_multinomial_logpmf(val, n, a) for val in vals]),
             np.asarray([model_single.fastlogp({"m": val}) for val in vals]),
             decimal=4,
         )
 
         assert_almost_equal(
-            np.asarray([dirichlet_multinomial_logpmf(v, n, a) for v in vals]),
+            np.asarray([dirichlet_multinomial_logpmf(val, n, a) for val in vals]),
             logpt(model_many.m, vals).eval().squeeze(),
             decimal=4,
         )
@@ -2319,7 +2325,7 @@ def test_dirichlet_multinomial_vec_1d_n(self):
         ns = np.array([10, 11])
 
         with Model() as model:
-            DirichletMultinomial("m", n=ns, a=a)
+            pm.DirichletMultinomial("m", n=ns, a=a)
 
         assert_almost_equal(
             sum(dirichlet_multinomial_logpmf(val, n, a) for val, n in zip(vals, ns)),
@@ -2355,7 +2361,6 @@ def test_dirichlet_multinomial_vec_2d_a(self):
             decimal=4,
         )
 
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_batch_dirichlet_multinomial(self):
         # Test that DM can handle a 3d array for `a`