Allow for batched alpha in StickBreakingWeights

pymc/distributions/multivariate.py

@@ -2192,33 +2192,24 @@ def make_node(self, rng, size, dtype, alpha, K):
         alpha = at.as_tensor_variable(alpha)
         K = at.as_tensor_variable(intX(K))
 
-        if alpha.ndim > 0:
-            raise ValueError("The concentration parameter needs to be a scalar.")
-
         if K.ndim > 0:
             raise ValueError("K must be a scalar.")
 
         return super().make_node(rng, size, dtype, alpha, K)
 
-    def _infer_shape(self, size, dist_params, param_shapes=None):
-        alpha, K = dist_params
-
-        size = tuple(size)
-
-        return size + (K + 1,)
+    def _supp_shape_from_params(self, dist_params, **kwargs):
+        K = dist_params[1]
+        return (K + 1,)
 
     @classmethod
     def rng_fn(cls, rng, alpha, K, size):
         if K < 0:
             raise ValueError("K needs to be positive.")
 
-        if size is None:
-            size = (K,)
-        elif isinstance(size, int):
-            size = (size,) + (K,)
-        else:
-            size = tuple(size) + (K,)
+        size = to_tuple(size)
+        size = np.broadcast_shapes(alpha.shape, size) + (K,)
 
+        alpha = alpha[..., np.newaxis]
         betas = rng.beta(1, alpha, size=size)
 
         sticks = np.concatenate(

pymc/tests/test_distributions.py

@@ -2291,6 +2291,24 @@ def test_dirichlet_multinomial_vectorized(self, n, a, extra_size):
                 3,
                 np.array([1.29317672, 1.50126157]),
             ),
+            (
+                np.array([5, 4, 3, 2, 1]) / 15,
+                np.array([0.5, 1, 2], dtype="float64"),
+                4,
+                np.array([1.51263013, 2.93119375, 2.99573227]),
+            ),
+            (
+                np.array([5, 4, 3, 2, 1]) / 15,
+                np.arange(1, 10, dtype="float64").reshape(3, 3),
+                4,
+                np.array(
+                    [
+                        [2.93119375, 2.99573227, 1.9095425],
+                        [0.35222059, -1.4632554, -3.44201938],
+                        [-5.53346686, -7.70739149, -9.94430955],
+                    ]
+                ),
+            ),
         ],
     )
     def test_stickbreakingweights_logp(self, value, alpha, K, logp):