Fix Dirichlet.logp by checking number of categories > 1 only at event dims

Author: Sayam753

pymc3/distributions/multivariate.py

@@ -522,7 +522,7 @@ def logp(self, value):
             tt.sum(logpow(value, a - 1) - gammaln(a), axis=-1) + gammaln(tt.sum(a, axis=-1)),
             tt.all(value >= 0),
             tt.all(value <= 1),
-            np.logical_not(a.broadcastable),
+            np.logical_not(a.broadcastable[-1]),
             tt.all(a > 0),
             broadcast_conditions=False,
         )

pymc3/tests/test_distributions.py

@@ -1696,6 +1696,11 @@ def test_lkj(self, x, eta, n, lp):
     def test_dirichlet(self, n):
         self.pymc3_matches_scipy(Dirichlet, Simplex(n), {"a": Vector(Rplus, n)}, dirichlet_logpdf)
 
+    def test_dirichlet_with_unit_batch_shape(self):
+        with pm.Model() as model:
+            a = pm.Dirichlet("a", a=np.ones((1, 2)))
+        np.isfinite(model.check_test_point()[0])
+
     def test_dirichlet_shape(self):
         a = tt.as_tensor_variable(np.r_[1, 2])
         with pytest.warns(DeprecationWarning):