Broadcast shapes of alpha and beta in Weibull rng (pymc-devs#7288)

mkusnetsov · Oct 26, 2024 · 2482aad · 2482aad
1 parent 278889f
commit 2482aad
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
@@ -2499,6 +2499,8 @@ def __call__(self, alpha, beta, size=None, **kwargs):
 
     @classmethod
     def rng_fn(cls, rng, alpha, beta, size) -> np.ndarray:
+        if size is None:
+            size = np.broadcast_shapes(alpha.shape, beta.shape)
         return np.asarray(beta * rng.weibull(alpha, size=size))
 
 

diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py
@@ -2369,6 +2369,14 @@ def seeded_weibul_rng_fn(self):
         "check_rv_size",
     ]
 
+    def test_rng_different_shapes(self):
+        # See issue #7220
+        rng = np.random.default_rng(123)
+        alpha = np.abs(rng.normal(size=5))
+        beta = np.abs(rng.normal(size=(3, 1)))
+        draws = pm.draw(pm.Weibull.dist(alpha, beta), random_seed=rng)
+        assert len(np.unique(draws)) == draws.size
+
 
 @pytest.mark.skipif(
     condition=_polyagamma_not_installed,