Improve support for dims in LKJCholeskyCov

pymc/distributions/multivariate.py

@@ -18,6 +18,7 @@
 import warnings
 
 from functools import reduce
+from itertools import product
 from typing import Optional
 
 import numpy as np
@@ -1433,14 +1434,32 @@
     """
 
     def __new__(cls, name, eta, n, sd_dist, *, compute_corr=True, store_in_trace=True, **kwargs):
+        dims = kwargs.pop("dims", None)
+
+        if dims is not None:
+            # TODO: Add check for 2d dims?
+            packed_dim_name, packed_dim_value = cls._make_packed_coord_from_dims(
+                n, dims, "packed_tril"
+            )
+            cls._register_new_coords_with_model(packed_dim_name, packed_dim_value)
+            kwargs["dims"] = [packed_dim_name]
+
         packed_chol = _LKJCholeskyCov(name, eta=eta, n=n, sd_dist=sd_dist, **kwargs)
+
         if not compute_corr:
             return packed_chol
         else:
             chol, corr, stds = cls.helper_deterministics(n, packed_chol)
             if store_in_trace:
-                corr = pm.Deterministic(f"{name}_corr", corr)
-                stds = pm.Deterministic(f"{name}_stds", stds)
+                corr_triu = corr[pt.triu_indices_from(corr, k=1)]
+                corr_triu_dim_name, corr_triu_dim_value = cls._make_packed_coord_from_dims(
+                    n, dims, "corr", lower=False, k=1
+                )
+                cls._register_new_coords_with_model(corr_triu_dim_name, corr_triu_dim_value)
+
+                corr_tril = pm.Deterministic(f"{name}_corr", corr_triu, dims=corr_triu_dim_name)
+                stds = pm.Deterministic(f"{name}_stds", stds, dims=dims[0])
+
             return chol, corr, stds
 
     @classmethod
@@ -1463,6 +1482,28 @@
         corr = inv_stds[None, :] * cov * inv_stds[:, None]
         return chol, corr, stds
 
+    @classmethod
+    def _make_packed_coord_from_dims(cls, n, dims, name_prefix, lower=True, k=0):
+        mod = pm.modelcontext(None)
+        chol_dims = [mod.coords[dim] for dim in dims]
+        if lower:
+            f_idx = np.tril_indices
+        else:
+            f_idx = np.triu_indices
+
+        flat_tri_idx = np.arange(n**2, dtype=int).reshape(n, n)[f_idx(n, k=k)]
+        coord_product = np.fromiter([f"{x}" for x in product(*chol_dims)], dtype="object")
+        tri_coords = coord_product[flat_tri_idx].tolist()
+
+        packed_dim_name = f"{name_prefix}_{dims[0]}"
+        return packed_dim_name, tri_coords
+
+    @classmethod
+    def _register_new_coords_with_model(cls, name, value):
+        mod = pm.modelcontext(None)
+        mod.coords[name] = value
+        mod.dim_lengths[name] = pt.TensorConstant(pt.lscalar, np.array(len(value)))
+
 
 class LKJCorrRV(RandomVariable):
     name = "lkjcorr"