Broadcast shapes of alpha and beta in Weibull rng

Author: tomicapretto

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))
 
 

tests/distributions/test_continuous.py

@@ -2369,6 +2369,25 @@ def seeded_weibul_rng_fn(self):
         "check_rv_size",
     ]
 
+    # See issue #7220
+    def test_rng_different_shapes(self):
+        rng = np.random.default_rng(123)
+
+        # Simulate mean from an only-intercept model. 2 chains, 100 draws, 5 observations.
+        # So 'mu' is the same for all the observations (because it's intercept-only)
+        mu_draws = np.abs(150 + np.dstack([rng.normal(size=(2, 100, 1))] * 5))
+
+        # Simulate some alpha values
+        alpha_draws = np.abs(rng.normal(size=(2, 100, 1)))
+
+        # With 'mu' and 'alpha' get 'beta', which is what pm.Weibull needs
+        beta_draws = mu_draws / sp.gamma(1 + 1 / alpha_draws)
+
+        # See the draws, for a given chain and draw, they look all the same!
+        weibull_draws = pm.draw(pm.Weibull.dist(alpha=alpha_draws, beta=beta_draws))
+
+        assert not (weibull_draws == weibull_draws[:, :, 0][..., None]).all()
+
 
 @pytest.mark.skipif(
     condition=_polyagamma_not_installed,