Allow for batched alpha in StickBreakingWeights

pymc/distributions/multivariate.py

@@ -2192,32 +2192,23 @@ 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) if size is not None else alpha.shape
+        size = size + (K,)
+        alpha = alpha[..., np.newaxis]
 
         betas = rng.beta(1, alpha, size=size)
 

pymc/tests/test_distributions.py

@@ -23,6 +23,7 @@
 from aeppl.logprob import ParameterValueError
 from aesara.tensor.random.utils import broadcast_params
 
+from pymc.aesaraf import compile_pymc
 from pymc.distributions.continuous import get_tau_sigma
 from pymc.util import UNSET
 
@@ -952,6 +953,15 @@ def test_hierarchical_obs_logp():
     assert not any(isinstance(o, RandomVariable) for o in ops)
 
 
+@pytest.fixture(scope="module")
+def _compile_stickbreakingweights_logpdf():
+    _value = at.vector()
+    _alpha = at.scalar()
+    _k = at.iscalar()
+    _logp = logp(StickBreakingWeights.dist(_alpha, _k), _value)
+    return compile_pymc([_value, _alpha, _k], _logp)
+
+
 class TestMatchesScipy:
     def test_uniform(self):
         check_logp(
@@ -2279,27 +2289,25 @@ def test_dirichlet_multinomial_vectorized(self, n, a, extra_size):
         )
 
     @pytest.mark.parametrize(
-        "value,alpha,K,logp",
+        "alpha,K",
         [
-            (np.array([5, 4, 3, 2, 1]) / 15, 0.5, 4, 1.5126301307277439),
-            (np.tile(1, 13) / 13, 2, 12, 13.980045245672827),
-            (np.array([0.001] * 10 + [0.99]), 0.1, 10, -22.971662448814723),
-            (np.append(0.5 ** np.arange(1, 20), 0.5**20), 5, 19, 94.20462772778092),
-            (
-                (np.array([[7, 5, 3, 2], [19, 17, 13, 11]]) / np.array([[17], [60]])),
-                2.5,
-                3,
-                np.array([1.29317672, 1.50126157]),
-            ),
+            (0.5, 4),
+            (2, 12),
+            (np.array([0.5, 1.0, 2.0]), 3),
+            (np.arange(1, 7, dtype="float64").reshape(2, 3), 5),
         ],
     )
-    def test_stickbreakingweights_logp(self, value, alpha, K, logp):
-        with Model() as model:
+    def test_stickbreakingweights_logp(self, alpha, K, _compile_stickbreakingweights_logpdf):
+        stickbreakingweights_logpdf = np.vectorize(
+            _compile_stickbreakingweights_logpdf, signature="(n),(),()->()"
+        )
+        value = pm.StickBreakingWeights.dist(alpha, K).eval()
+        with Model():
             sbw = StickBreakingWeights("sbw", alpha=alpha, K=K, transform=None)
         pt = {"sbw": value}
         assert_almost_equal(
             pm.logp(sbw, value).eval(),
-            logp,
+            stickbreakingweights_logpdf(value, alpha, K),
             decimal=select_by_precision(float64=6, float32=2),
             err_msg=str(pt),
         )

pymc/tests/test_distributions_random.py

@@ -1321,6 +1321,18 @@ def check_basic_properties(self):
         assert np.all(draws <= 1)
 
 
+class TestStickBreakingWeights_1D_alpha(BaseTestDistributionRandom):
+    pymc_dist = pm.StickBreakingWeights
+    pymc_dist_params = {"alpha": [1.0, 2.0, 3.0], "K": 19}
+    expected_rv_op_params = {"alpha": [1.0, 2.0, 3.0], "K": 19}
+    sizes_to_check = [None, (3,), (5, 3)]
+    sizes_expected = [(3, 20), (3, 20), (5, 3, 20)]
+    checks_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_rv_size",
+    ]
+
+
 class TestCategorical(BaseTestDistributionRandom):
     pymc_dist = pm.Categorical
     pymc_dist_params = {"p": np.array([0.28, 0.62, 0.10])}