Repository: pyro-ppl/pyro

Issue:
Difference in event_shape between pyro's LKJCholesky and torch's
### Issue Description
Event shape of inverse transform of `LKJCholesky` seems to differ from torch's, and of samples from distribution.

### Environment
 - WSL2 Ubuntu 20.04
 - Python 3.9
 - torch 1.13.1
 - pyro 1.8.4+9ed468d

### Code Snippet
```python
from torch.distributions import TransformedDistribution, LKJCholesky
from torch.distributions import biject_to

# torch
torch_lkj = LKJCholesky(2, 0.5)
torch_inv = biject_to(torch_lkj.support).inv
torch_transformed = TransformedDistribution(torch_lkj, torch_inv)

# pyro
from pyro.distributions import LKJCholesky as PyroCholesky
from pyro import __version__

print(__version__)

pyro_lkj = PyroCholesky(2, 0.5)
pyro_inv = biject_to(pyro_lkj.support).inv
pyro_transformed = TransformedDistribution(pyro_lkj, pyro_inv)

assert pyro_transformed.event_shape == torch_transformed.event_shape
```
Stepping through the code, it seems as though `pyro.distributions.transforms.cholesky.CorrLCholeskyTransform` doesn't implement the `inverse_shape` method, whereas torch does with `torch.distributions.transforms.CorrCholeskyTransform`. This results in pyro returning an event shape corresponding to `torch.Size([2])`, wheras torch returns `torch.Size([1])`. 



Discussion:
It looks like `pyro.distributions.transforms.cholesky.CorrLCholeskyTransform` should be deprecated in favor of `torch.distributions.transforms.CorrCholeskyTransform`. Is that true @fritzo @fehiepsi?
Yes, I think we can use the upstream one.

Resolving patch:
diff --git a/pyro/distributions/transforms/__init__.py b/pyro/distributions/transforms/__init__.py
--- a/pyro/distributions/transforms/__init__.py
+++ b/pyro/distributions/transforms/__init__.py
@@ -32,6 +32,7 @@
 from .block_autoregressive import BlockAutoregressive, block_autoregressive
 from .cholesky import (
     CholeskyTransform,
+    CorrCholeskyTransform,
     CorrLCholeskyTransform,
     CorrMatrixCholeskyTransform,
 )
@@ -90,17 +91,11 @@ def _transform_to_sphere(constraint):
     return Normalize()
 
 
-@biject_to.register(constraints.corr_cholesky)
-@transform_to.register(constraints.corr_cholesky)
-def _transform_to_corr_cholesky(constraint):
-    return CorrLCholeskyTransform()
-
-
 @biject_to.register(constraints.corr_matrix)
 @transform_to.register(constraints.corr_matrix)
 def _transform_to_corr_matrix(constraint):
     return ComposeTransform(
-        [CorrLCholeskyTransform(), CorrMatrixCholeskyTransform().inv]
+        [CorrCholeskyTransform(), CorrMatrixCholeskyTransform().inv]
     )
 
 
diff --git a/pyro/distributions/transforms/cholesky.py b/pyro/distributions/transforms/cholesky.py
--- a/pyro/distributions/transforms/cholesky.py
+++ b/pyro/distributions/transforms/cholesky.py
@@ -2,89 +2,21 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
+import warnings
 
 import torch
-from torch.distributions.transforms import Transform
+from torch.distributions.transforms import CorrCholeskyTransform, Transform
 
 from .. import constraints
 
 
-def _vector_to_l_cholesky(z):
-    D = (1.0 + math.sqrt(1.0 + 8.0 * z.shape[-1])) / 2.0
-    if D % 1 != 0:
-        raise ValueError("Correlation matrix transformation requires d choose 2 inputs")
-    D = int(D)
-    x = torch.zeros(z.shape[:-1] + (D, D), dtype=z.dtype, device=z.device)
-
-    x[..., 0, 0] = 1
-    x[..., 1:, 0] = z[..., : (D - 1)]
-    i = D - 1
-    last_squared_x = torch.zeros(z.shape[:-1] + (D,), dtype=z.dtype, device=z.device)
-    for j in range(1, D):
-        distance_to_copy = D - 1 - j
-        last_squared_x = last_squared_x[..., 1:] + x[..., j:, (j - 1)].clone() ** 2
-        x[..., j, j] = (1 - last_squared_x[..., 0]).sqrt()
-        x[..., (j + 1) :, j] = (
-            z[..., i : (i + distance_to_copy)] * (1 - last_squared_x[..., 1:]).sqrt()
-        )
-        i += distance_to_copy
-    return x
-
-
-class CorrLCholeskyTransform(Transform):
-    """
-    Transforms a vector into the cholesky factor of a correlation matrix.
-
-    The input should have shape `[batch_shape] + [d * (d-1)/2]`. The output will
-    have shape `[batch_shape] + [d, d]`.
-
-    References:
-
-    [1] Cholesky Factors of Correlation Matrices. Stan Reference Manual v2.18,
-    Section 10.12.
-
-    """
-
-    domain = constraints.real_vector
-    codomain = constraints.corr_cholesky
-    bijective = True
-
-    def __eq__(self, other):
-        return isinstance(other, CorrLCholeskyTransform)
-
-    def _call(self, x):
-        z = x.tanh()
-        return _vector_to_l_cholesky(z)
-
-    def _inverse(self, y):
-        if y.shape[-2] != y.shape[-1]:
-            raise ValueError(
-                "A matrix that isn't square can't be a Cholesky factor of a correlation matrix"
-            )
-        D = y.shape[-1]
-
-        z_tri = torch.zeros(
-            y.shape[:-2] + (D - 2, D - 2), dtype=y.dtype, device=y.device
-        )
-        z_stack = [y[..., 1:, 0]]
-
-        for i in range(2, D):
-            z_tri[..., i - 2, 0 : (i - 1)] = (
-                y[..., i, 1:i] / (1 - y[..., i, 0 : (i - 1)].pow(2).cumsum(-1)).sqrt()
-            )
-        for j in range(D - 2):
-            z_stack.append(z_tri[..., j:, j])
-
-        z = torch.cat(z_stack, -1)
-        return torch.log1p((2 * z) / (1 - z)) / 2
-
-    def log_abs_det_jacobian(self, x, y):
-        # Note dependence on pytorch 1.0.1 for batched tril
-        tanpart = x.cosh().log().sum(-1).mul(-2)
-        matpart = (
-            (1 - y.pow(2).cumsum(-1).tril(diagonal=-2)).log().div(2).sum(-1).sum(-1)
+class CorrLCholeskyTransform(CorrCholeskyTransform):  # DEPRECATED
+    def __init__(self, cache_size=0):
+        warnings.warn(
+            "class CorrLCholeskyTransform is deprecated in favor of CorrCholeskyTransform.",
+            FutureWarning,
         )
-        return tanpart + matpart
+        super().__init__(cache_size=cache_size)
 
 
 class CholeskyTransform(Transform):


Test patch:
diff --git a/tests/distributions/test_lkj.py b/tests/distributions/test_lkj.py
--- a/tests/distributions/test_lkj.py
+++ b/tests/distributions/test_lkj.py
@@ -39,7 +39,7 @@ def _autograd_log_det(ys, x):
 
 @pytest.mark.parametrize("y_shape", [(1,), (3, 1), (6,), (1, 6), (2, 6)])
 def test_unconstrained_to_corr_cholesky_transform(y_shape):
-    transform = transforms.CorrLCholeskyTransform()
+    transform = transforms.CorrCholeskyTransform()
     y = torch.empty(y_shape).uniform_(-4, 4).requires_grad_()
     x = transform(y)
 
diff --git a/tests/distributions/test_transforms.py b/tests/distributions/test_transforms.py
--- a/tests/distributions/test_transforms.py
+++ b/tests/distributions/test_transforms.py
@@ -418,7 +418,7 @@ def test_cholesky_transform(batch_shape, dim, transform):
         tril_mask = arange < arange.view(-1, 1)
     else:
         tril_mask = arange < arange.view(-1, 1) + 1
-    x = transform.inv(T.CorrLCholeskyTransform()(z))  # creates corr_matrix
+    x = transform.inv(T.CorrCholeskyTransform()(z))  # creates corr_matrix
 
     def vec_to_mat(x_vec):
         x_mat = x_vec.new_zeros(batch_shape + (dim, dim))