Rename transformer to CholeskyCorrTransformer

Author: jessegrabowski

pymc/distributions/multivariate.py

@@ -74,7 +74,7 @@
     to_tuple,
 )
 from pymc.distributions.transforms import (
-    CholeskyCorr,
+    CholeskyCorrTransform,
     Interval,
     ZeroSumTransform,
     _default_transform,
@@ -1652,9 +1652,9 @@ def logp(value, n, eta):
 
 @_default_transform.register(_LKJCorr)
 def lkjcorr_default_transform(op, rv):
-    _, _, _, n, *_ = rv.owner.inputs
+    rng, shape, n, eta, *_ = rv.owner.inputs = rv.owner.inputs
     n = pt.get_scalar_constant_value(n)  # Safely extract scalar value without eval
-    return CholeskyCorr(n)
+    return CholeskyCorrTransform(n=n, upper=False)
 
 
 class LKJCorr:

pymc/distributions/transforms.py

@@ -34,7 +34,7 @@
 __all__ = [
     "Chain",
     "Chain",
-    "CholeskyCorr",
+    "CholeskyCorrTransform",
     "CholeskyCovPacked",
     "CholeskyCovPacked",
     "Interval",
@@ -140,7 +140,7 @@ def log_jac_det(self, value, *inputs):
         return pt.sum(y, axis=-1)
 
 
-class CholeskyCorr(Transform):
+class CholeskyCorrTransform(Transform):
     """
     Map an unconstrained real vector the Cholesky factor of a correlation matrix.
 
@@ -181,7 +181,7 @@ class CholeskyCorr(Transform):
        https://github.com/tensorflow/probability/
     """
 
-    name = "cholesky-corr"
+    name = "cholesky_corr"
 
     def __init__(self, n, upper: bool = False):
         """
@@ -267,15 +267,14 @@ def _fill_triangular_spiral(
         upper = self.upper
 
         if unit_diag:
-            m -= n
-            n -= 1
+            n = n - 1
 
         tail = x_raveled[..., n:]
 
         if upper:
-            xc = pt.concatenate([x_raveled, pt.flip(tail, -1)])
+            xc = pt.concatenate([x_raveled, pt.flip(tail, -1)], axis=-1)
         else:
-            xc = pt.concatenate([tail, pt.flip(x_raveled, -1)])
+            xc = pt.concatenate([tail, pt.flip(x_raveled, -1)], axis=-1)
 
         y = pt.reshape(xc, (*batch_shape, n, n))
         return pt.triu(y) if upper else pt.tril(y)
@@ -306,8 +305,8 @@ def _inverse_fill_triangular_spiral(
         n, m = self.n, self.m
 
         if unit_diag:
-            m -= n
-            n -= 1
+            m = m - n
+            n = n - 1
 
         upper = self.upper
 

tests/distributions/test_transform.py

@@ -667,7 +667,7 @@ def log_jac_det(self, value, *inputs):
             m.logp(jacobian=jacobian_val)
 
 
-class TestLJKCholeskyCorr:
+class TestLJKCholeskyCorrTransform:
     def _get_test_values(self):
         x_unconstrained = np.array([2.0, 2.0, 1.0])
         x_constrained = np.array(
@@ -696,7 +696,7 @@ def test_fill_triangular_spiral(self, upper):
                 ]
             )
 
-        transform = tr.CholeskyCorr(n=3, upper=upper)
+        transform = tr.CholeskyCorrTransform(n=3, upper=upper)
 
         np.testing.assert_allclose(
             transform._fill_triangular_spiral(x_unconstrained, unit_diag=False).eval(),
@@ -709,7 +709,7 @@ def test_fill_triangular_spiral(self, upper):
         )
 
     def test_forward(self):
-        transform = tr.CholeskyCorr(n=3, upper=False)
+        transform = tr.CholeskyCorrTransform(n=3, upper=False)
         x_unconstrained, x_constrained = self._get_test_values()
 
         np.testing.assert_allclose(
@@ -719,7 +719,7 @@ def test_forward(self):
         )
 
     def test_backward(self):
-        transform = tr.CholeskyCorr(n=3, upper=False)
+        transform = tr.CholeskyCorrTransform(n=3, upper=False)
         x_unconstrained, x_constrained = self._get_test_values()
 
         np.testing.assert_allclose(
@@ -729,7 +729,7 @@ def test_backward(self):
         )
 
     def test_transform_round_trip(self):
-        transform = tr.CholeskyCorr(n=3, upper=False)
+        transform = tr.CholeskyCorrTransform(n=3, upper=False)
         x_unconstrained, x_constrained = self._get_test_values()
 
         constrained_reconstructed = transform.backward(transform.forward(x_constrained)).eval()
@@ -739,7 +739,7 @@ def test_transform_round_trip(self):
         np.testing.assert_allclose(x_constrained, constrained_reconstructed, atol=1e-6)
 
     def test_log_jac_det(self):
-        transform = tr.CholeskyCorr(n=3, upper=False)
+        transform = tr.CholeskyCorrTransform(n=3, upper=False)
         x_unconstrained, x_constrained = self._get_test_values()
 
         computed_log_jac_det = transform.log_jac_det(x_unconstrained).eval()