Updated LKJCholeskyCovRV to use univariate distributions only.

pymc/distributions/multivariate.py

@@ -1061,8 +1061,6 @@ def rng_fn(self, rng, eta, n, sd_dist, size=None):
             size = 1
             orig_size = None
 
-        # TODO: This seems to be more or less the same code inside LKJCorrRV.rng_fn,
-        #  Can't we just call it?
         P = np.eye(n) * np.ones((size,) + (n, n))
         # original implementation in R see:
         # https://github.com/rmcelreath/rethinking/blob/master/R/distributions.r
@@ -1089,34 +1087,9 @@ def rng_fn(self, rng, eta, n, sd_dist, size=None):
         for _var in sd_dist_inputs[3:]:
             rv_inputs.append(_var.value)
 
-        # TODO: We can inspect the ndims_params of the sd_dist op to know what the
-        #  dimensions are, without doing that first call to rng_fn
-        #  Perhaps we can simply limit sd_dists to be univariate distributions, which
-        #  should simplify the logic quite a lot
+        # Here we simply limit sd_dists to be univariate distributions
         rv_inputs.append(P.shape[:-2])
-        D = np.atleast_1d(sd_dist_op.rng_fn(*rv_inputs))
-        if D.shape in [tuple(), (1,)]:
-            rv_inputs[-1] = P.shape[:-1]
-            D = sd_dist_op.rng_fn(*rv_inputs)
-        elif D.ndim < C.ndim - 1:
-            D = [D] + [sd_dist_op.rng_fn(*rv_inputs) for _ in range(n - 1)]
-            D = np.moveaxis(np.array(D), 0, C.ndim - 2)
-        elif D.ndim == C.ndim - 1:
-            if D.shape[-1] == 1:
-                D = [D] + [sd_dist_op.rng_fn(*rv_inputs) for _ in range(n - 1)]
-                D = np.concatenate(D, axis=-1)
-            elif D.shape[-1] != n:
-                raise ValueError(
-                    "The size of the samples drawn from the "
-                    "supplied sd_dist.random have the wrong "
-                    "size. Expected {} but got {} instead.".format(n, D.shape[-1])
-                )
-        else:
-            raise ValueError(
-                "Supplied sd_dist.random generates samples with "
-                "too many dimensions. It must yield samples "
-                "with 0 or 1 dimensions. Got {} instead".format(D.ndim - C.ndim - 2)
-            )
+        D = sd_dist_op.rng_fn(*rv_inputs)
 
         C *= D[..., :, np.newaxis] * D[..., np.newaxis, :]
         tril_idx = np.tril_indices(n, k=0)
@@ -1142,6 +1115,9 @@ def __new__(cls, *args, **kwargs):
             kwargs["transform"] = cls.default_transform()
 
         sd_dist = kwargs["sd_dist"]
+        assert (
+            sd_dist.ndim == 0
+        ), "LKJCholeskyCov only supports univariate distributions as distribution for the standard deviations."
 
         rv_op = type(
             f"_lkjcholeskycov",
@@ -1163,16 +1139,16 @@ def __new__(cls, *args, **kwargs):
         def logp(op, value_var_list, *dist_params, **kwargs):
             _dist_params = dist_params[3:]
             value_var = value_var_list[0]
-            return cls.logp(value_var, *dist_params, **kwargs)
+            return cls.logp(value_var, *_dist_params)
 
         cls.rv_op = rv_op
 
         return super().__new__(cls, *args, **kwargs)
 
     @classmethod
     def default_transform(cls):
-        def transform_params(rv_var):
-            _, _, _, eta, n, _ = rv_var.owner.inputs
+        def transform_params(rv_inputs):
+            _, _, _, eta, n, _ = rv_inputs
             return np.arange(1, n.data + 1).cumsum() - 1
 
         return transforms.CholeskyCovPacked(transform_params)
@@ -1500,7 +1476,7 @@ def logp(x, n, eta):
         tri_index[np.triu_indices(_n, k=1)] = np.arange(shape)
         tri_index[np.triu_indices(_n, k=1)[::-1]] = np.arange(shape)
 
-        X = at.take(x, tri_index)
+        X = at.subtensor.advanced_subtensor(x, tri_index)
         X = at.fill_diagonal(X, 1)
 
         result = _lkj_normalizing_constant(_eta, _n)

pymc/tests/test_distributions.py

@@ -3361,4 +3361,7 @@ def test_lkjcholeskycov():
     pt = np.random.random(120)
     pt = {"packedL_cholesky-cov-packed__": pt}
     logp = model.fastlogp(pt)
-    assert 0  # Test not complete
+    assert logp is not None
+
+
+test_lkjcholeskycov()