Automatically resize sd_dist in _LKJCholeskyCov

Author: ricardoV94

pymc/distributions/multivariate.py

@@ -41,7 +41,7 @@
 
 import pymc as pm
 
-from pymc.aesaraf import floatX, intX
+from pymc.aesaraf import change_rv_size, floatX, intX
 from pymc.distributions import transforms
 from pymc.distributions.continuous import (
     BoundedContinuous,
@@ -1199,8 +1199,21 @@ def dist(cls, eta, n, sd_dist, size=None, **kwargs):
             isinstance(sd_dist, Variable)
             and sd_dist.owner is not None
             and isinstance(sd_dist.owner.op, RandomVariable)
+            and sd_dist.owner.op.ndim_supp < 2
         ):
-            raise TypeError("sd_dist must be a Distribution variable")
+            raise TypeError("sd_dist must be a scalar or vector distribution variable")
+
+        # We resize the sd_dist automatically so that it has (size x n) independent draws
+        # which is what the `_LKJCholeskyCovRV.rng_fn` expects. This makes the random
+        # and logp methods equivalent, as the latter also assumes a unique value for each
+        # diagonal element.
+        # Since `eta` and `n` are forced to be scalars we don't need to worry about
+        # implied batched dimensions for the time being.
+        if sd_dist.owner.op.ndim_supp == 0:
+            sd_dist = change_rv_size(sd_dist, to_tuple(size) + (n,))
+        else:
+            # The support shape must be `n` but we have no way of controlling it
+            sd_dist = change_rv_size(sd_dist, to_tuple(size))
 
         # sd_dist is part of the generative graph, but should be completely ignored
         # by the logp graph, since the LKJ logp explicitly includes these terms.
@@ -1288,7 +1301,9 @@ class LKJCholeskyCov:
     n: int
         Dimension of the covariance matrix (n > 1).
     sd_dist: pm.Distribution
-        A distribution for the standard deviations, should have `size=n`.
+        A positive scalar or vector distribution for the standard deviations, created
+        with the `.dist()` API. Should have `shape[-1]=n`. Scalar distributions will be
+        automatically resized to ensure this.
     compute_corr: bool, default=True
         If `True`, returns three values: the Cholesky decomposition, the correlations
         and the standard deviations of the covariance matrix. Otherwise, only returns

pymc/tests/test_distributions.py

@@ -3382,7 +3382,7 @@ def test_dist(self):
     def test_sd_dist_distribution(self):
         with pm.Model() as m:
             sd_dist = at.constant([1, 2, 3])
-            with pytest.raises(TypeError, match="sd_dist must be a Distribution variable"):
+            with pytest.raises(TypeError, match="^sd_dist must be a scalar or vector distribution"):
                 x = pm.LKJCholeskyCov("x", n=3, eta=1, sd_dist=sd_dist)
 
     def test_sd_dist_registered(self):
@@ -3402,3 +3402,17 @@ def test_no_warning_logp(self):
         with pytest.warns(None) as record:
             m.logpt()
         assert not record
+
+    @pytest.mark.parametrize(
+        "sd_dist",
+        [
+            pm.Exponential.dist(1),
+            pm.MvNormal.dist(np.ones(3), np.eye(3)),
+        ],
+    )
+    def test_sd_dist_automatically_resized(self, sd_dist):
+        x = pm.LKJCholeskyCov.dist(n=3, eta=1, sd_dist=sd_dist, size=10, compute_corr=False)
+        resized_sd_dist = x.owner.inputs[-1]
+        assert resized_sd_dist.eval().shape == (10, 3)
+        # LKJCov has support shape `(n * (n+1)) // 2`
+        assert x.eval().shape == (10, 6)