Fix failing default transform for LKJCorr

pymc/logprob/transform_value.py

@@ -125,13 +125,12 @@ def transformed_value_logprob(op, values, *rv_outs, use_jacobian=True, **kwargs)
             raise NotImplementedError(
                 f"Univariate transform {transform} cannot be applied to multivariate {rv_op}"
             )
-        else:
-            # Check there is no broadcasting between logp and jacobian
-            if logp.type.broadcastable != log_jac_det.type.broadcastable:
-                raise ValueError(
-                    f"The logp of {rv_op} and log_jac_det of {transform} are not allowed to broadcast together. "
-                    "There is a bug in the implementation of either one."
-                )
+        # Check there is no broadcasting between logp and jacobian
+        if logp.type.broadcastable != log_jac_det.type.broadcastable:
+            raise ValueError(
+                f"The logp of {rv_op} and log_jac_det of {transform} are not allowed to broadcast together. "
+                "There is a bug in the implementation of either one."
+            )
 
         if use_jacobian:
             if value.name:

tests/distributions/test_multivariate.py

@@ -13,7 +13,6 @@
 #   limitations under the License.
 
 import functools as ft
-import re
 import warnings
 
 import numpy as np
@@ -2122,11 +2121,25 @@ def ref_rand(size, n, eta):
             size=1000,
         )
 
-    def test_default_transform(self):
+    @pytest.mark.parametrize(
+        argnames="shape, expected_shape",
+        argvalues=[
+            ((2,), ()),
+            pytest.param(
+                (3, 2),
+                (3,),
+                marks=pytest.mark.xfail(
+                    raises=NotImplementedError,
+                    reason="We do not support batch dimensions for pm.LKJCorr yet.",
+                ),
+            ),
+        ],
+    )
+    def test_default_transform(self, shape, expected_shape):
         with pm.Model() as m:
-            x = pm.LKJCorr("x", n=2, eta=1, shape=(3, 2))
+            x = pm.LKJCorr("x", n=2, eta=1, shape=shape)
         assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform)
-        assert m.logp(sum=False)[0].shape == (3,)
+        assert m.logp(sum=False)[0].type.shape == expected_shape
 
 
 class TestLKJCholeskyCov(BaseTestDistributionRandom):