Repository: pyro-ppl/numpyro

Issue:
Incorrect event dimension for `ComposeTransform` if transform reduces event dimensions
An example is probably the clearest illustration. If the transform increases the number of event dimensions, everything works as expected.

```python
>>> from numpyro.distributions.transforms import CorrCholeskyTransform, ComposeTransform
>>> part = CorrCholeskyTransform()
>>> part.domain.event_dim, part.codomain.event_dim
(1, 2)
>>> composed = ComposeTransform([part])
>>> composed.domain.event_dim, composed.codomain.event_dim
(1, 2)
```

If the transform reduces the number of event dimensions, wrapping the transform in a `ComposeTransform` leads to unexpected results.

```python
>>> part = CorrCholeskyTransform().inv
>>> part.domain.event_dim, part.codomain.event_dim
(2, 1)
>>> composed = ComposeTransform([part])
>>> composed.domain.event_dim, composed.codomain.event_dim
(3, 1)
```

I had a brief look at the code below, but I couldn't quite get my head around it.

https://github.com/pyro-ppl/numpyro/blob/8ace34f1aca094897328730853b021d53b230f65/numpyro/distributions/transforms.py#L292-L298

Here's a minimal test that can be added to `test/test_transforms.py` to reproduce.

```python
@pytest.mark.parametrize("transform", [
    CorrCholeskyTransform(),  # passes
    CorrCholeskyTransform().inv,  # fails
])
def test_compose_domain_codomain(transform):
    composed = ComposeTransform([transform])
    assert transform.domain.event_dim == composed.domain.event_dim
    assert transform.codomain.event_dim == composed.codomain.event_dim
```

Discussion:

Patch:
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
index 290e504c..05ae19ef 100644
--- a/numpyro/distributions/transforms.py
+++ b/numpyro/distributions/transforms.py
@@ -282,7 +282,7 @@ class AffineTransform(Transform):
 
 def _get_compose_transform_input_event_dim(parts):
     input_event_dim = parts[-1].domain.event_dim
-    for part in parts[len(parts) - 1 :: -1]:
+    for part in parts[:-1][::-1]:
         input_event_dim = part.domain.event_dim + max(
             input_event_dim - part.codomain.event_dim, 0
         )


Test patch:
diff --git a/test/test_transforms.py b/test/test_transforms.py
index cecbf2ca..9d8c6594 100644
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -397,3 +397,16 @@ def test_biject_to(constraint, shape):
     expected_shape = constrained.shape[: constrained.ndim - constraint.event_dim]
     assert passed.shape == expected_shape
     assert jnp.all(passed)
+
+
+@pytest.mark.parametrize(
+    "transform",
+    [
+        CorrCholeskyTransform(),
+        CorrCholeskyTransform().inv,
+    ],
+)
+def test_compose_domain_codomain(transform):
+    composed = ComposeTransform([transform])
+    assert transform.domain.event_dim == composed.domain.event_dim
+    assert transform.codomain.event_dim == composed.codomain.event_dim