Tweak DirichletMultinomial logp and refactor some multivariate logp tests

conda-envs/environment-dev-py37.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-dev-py38.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-dev-py39.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py37.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py38.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py39.yml

@@ -4,7 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools

conda-envs/windows-environment-dev-py38.yml

@@ -4,7 +4,7 @@ channels:
 - defaults
 dependencies:
  # base dependencies (see install guide for Windows)
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/windows-environment-test-py38.yml

@@ -4,7 +4,7 @@ channels:
 - defaults
 dependencies:
  # base dependencies (see install guide for Windows)
-- aeppl>=0.0.17
+- aeppl=0.0.17
 - aesara>=2.2.6
 - arviz>=0.11.2
 - cachetools

pymc/distributions/multivariate.py

@@ -406,8 +406,9 @@ class Dirichlet(Continuous):
 
     Parameters
     ----------
-    a: array
-        Concentration parameters (a > 0).
+    a: float array
+        Concentration parameters (a > 0). The number of categories is given by the
+        length of the last axis.
     """
     rv_op = dirichlet
 
@@ -507,13 +508,12 @@ class Multinomial(Discrete):
 
     Parameters
     ----------
-    n: int or array
-        Number of trials (n > 0). If n is an array its shape must be (N,) with
-        N = p.shape[0]
-    p: one- or two-dimensional array
-        Probability of each one of the different outcomes. Elements must
-        be non-negative and sum to 1 along the last axis. They will be
-        automatically rescaled otherwise.
+    n: int
+        Total counts in each replicate (n > 0).
+    p: float array
+        Probability of each one of the different outcomes (0 <= p <= 1). The number of
+        categories is given by the length of the last axis. Elements are expected to sum
+        to 1 along the last axis, and they will be automatically rescaled otherwise.
     """
     rv_op = multinomial
 
@@ -626,17 +626,12 @@ class DirichletMultinomial(Discrete):
 
     Parameters
     ----------
-    n : int or array
-        Total counts in each replicate. If n is an array its shape must be (N,)
-        with N = a.shape[0]
-
-    a : one- or two-dimensional array
-        Dirichlet parameter. Elements must be strictly positive.
-        The number of categories is given by the length of the last axis.
+    n : int
+        Total counts in each replicate (n > 0).
 
-    shape : integer tuple
-        Describes shape of distribution. For example if n=array([5, 10]), and
-        a=array([1, 1, 1]), shape should be (2, 3).
+    a : float array
+        Dirichlet concentration parameters (a > 0). The number of categories is given by
+        the length of the last axis.
     """
     rv_op = dirichlet_multinomial
 
@@ -661,15 +656,10 @@ def logp(value, n, a):
         -------
         TensorVariable
         """
-        if value.ndim >= 1:
-            n = at.shape_padright(n)
-            if a.ndim > 1:
-                a = at.shape_padleft(a)
-
-        sum_a = a.sum(axis=-1, keepdims=True)
+        sum_a = a.sum(axis=-1)
         const = (gammaln(n + 1) + gammaln(sum_a)) - gammaln(n + sum_a)
         series = gammaln(value + a) - (gammaln(value + 1) + gammaln(a))
-        result = const + series.sum(axis=-1, keepdims=True)
+        result = const + series.sum(axis=-1)
 
         # Bounds checking to confirm parameters and data meet all constraints
         # and that each observation value_i sums to n_i.
@@ -678,13 +668,10 @@ def logp(value, n, a):
             value >= 0,
             a > 0,
             n >= 0,
-            at.eq(value.sum(axis=-1, keepdims=True), n),
+            at.eq(value.sum(axis=-1), n),
             broadcast_conditions=False,
         )
 
-    def _distr_parameters_for_repr(self):
-        return ["n", "a"]
-
 
 class _OrderedMultinomial(Multinomial):
     r"""