Fix ZeroSumTransform logp for non-trailing zero-sum axes

pymc/distributions/transforms.py

@@ -683,6 +683,11 @@ def backward(self, value, *rv_inputs):
             value = self.extend_axis(value, axis=axis)
         return value
 
+    # Zero-summed axes, used by `transformed_value_logprob` to reduce the matching logp axes.
+    @property
+    def jacobian_reduce_axes(self):
+        return self.zerosum_axes
+
     def log_jac_det(self, value, *rv_inputs):
         return value.sum(self.zerosum_axes).zeros_like()
 

pymc/logprob/transform_value.py

@@ -102,10 +102,14 @@ def transformed_value_logprob(op, values, *rv_outs, use_jacobian=True, **kwargs)
         log_jac_det = transform.log_jac_det(original_forward_value, *rv_inputs).copy()
         # The jacobian determinant has less dims than the logp
         # when a multivariate transform (like Simplex or Ordered) is applied to univariate distributions.
-        # In this case we have to reduce the last logp dimensions, as they are no longer independent
+        # In this case we have to reduce those logp dimensions, as they are no longer independent.
+        # Default to the trailing axes; transforms may override which axes via `jacobian_reduce_axes`.
         if log_jac_det.ndim < logp.ndim:
             diff_ndims = logp.ndim - log_jac_det.ndim
-            logp = logp.sum(axis=np.arange(-diff_ndims, 0))
+            reduce_axes = getattr(transform, "jacobian_reduce_axes", None)
+            if reduce_axes is None:
+                reduce_axes = np.arange(-diff_ndims, 0)
+            logp = logp.sum(axis=tuple(reduce_axes))
         # This case is sometimes, but not always, trivial to accommodate depending on the "space rank" of the
         # multivariate distribution. See https://proceedings.mlr.press/v130/radul21a.html
         elif log_jac_det.ndim > logp.ndim:

tests/distributions/test_multivariate.py

@@ -1780,6 +1780,26 @@ def test_batched_transformed_logp_shape(self):
         assert m.logp(sum=False)[0].type.shape == (3,)
         assert m.logp(sum=False, jacobian=False)[0].type.shape == (3,)
 
+    def test_zerosum_transform_non_trailing_axis(self):
+        # Regression: zero-summing a non-trailing axis must reduce the logp over that axis, else
+        # logp and jacobian shapes mismatch once dims are static (e.g. after freeze_dims_and_data).
+        from pymc.distributions.transforms import ZeroSumTransform
+        from pymc.model.transform.optimization import freeze_dims_and_data
+
+        # Static shape: reduced over the zero-summed (leading) axis, keeping the trailing axis.
+        with pm.Model() as m_static:
+            pm.Normal("x", 0.0, 1.0, shape=(3, 2), transform=ZeroSumTransform([0]))
+        assert m_static.logp(sum=False)[0].type.shape == (2,)
+
+        # Dynamic dims build fine; freezing makes shapes static and must not raise.
+        with pm.Model(coords={"a": range(3), "b": range(2)}) as m:
+            pm.Normal("x", 0.0, 1.0, dims=("a", "b"), transform=ZeroSumTransform([0]))
+        frozen = freeze_dims_and_data(m)
+        np.testing.assert_allclose(
+            m.compile_logp()(m.initial_point()),
+            frozen.compile_logp()(frozen.initial_point()),
+        )
+
 
 class TestMvStudentTCov(BaseTestDistributionRandom):
     def mvstudentt_rng_fn(self, size, nu, mu, scale, rng):