BUG: ZeroSumTransform jacobian shape breaks on model freeze

Describe the issue:

ZeroSumTransform.log_jac_det returns a shaped zero::

def log_jac_det(self, value, *rv_inputs):
    return value.sum(self.zerosum_axes).zeros_like()

i.e. zeros with the axes that were NOT zero-summed (here the trailing axis, size 2).

But pymc/logprob/transform_value.py::transformed_value_logprob, when the jacobian has
fewer dims than the logp, reduces the logp's trailing axes to match::

if log_jac_det.ndim < logp.ndim:
    diff_ndims = logp.ndim - log_jac_det.ndim
    logp = logp.sum(axis=np.arange(-diff_ndims, 0))   # keeps the LEADING axes

So for zerosum_axes=[0] on a (3, 2) variable:

• log_jac_det keeps the trailing axis -> shape (2,)
• logp is reduced over trailing axes -> shape (3,)
• logp + log_jac_det -> (3,) + (2,) -> incompatible.

With symbolic shapes (dims, no freezing) both reduce to (None,) and pytensor lets them
"broadcast", so the bug is silent. With static shapes -- e.g. an explicit shape=, or
after pymc.model.transform.optimization.freeze_dims_and_data (which nutpie's numba/jax
backends apply) -- the shapes are concrete and graph construction raises.

ZeroSumNormal itself is unaffected because it only zero-sums the trailing n axes, where
the kept (leading) axes coincidentally match pymc's reduction. The bug shows up when
ZeroSumTransform is applied directly to a non-trailing axis (a supported public transform,
e.g. to zero-sum a leading/categorical axis of a multi-dim effect).

Reproduceable code example:

import pymc as pm
import pytensor
from pymc.distributions.transforms import ZeroSumTransform

print("pymc:", pm.__version__, "| pytensor:", pytensor.__version__)


def build_logp(zerosum_axes):
    with pm.Model() as m:
        # static, non-square shape so the bug is not hidden behind symbolic dims
        pm.Normal("x", 0.0, 1.0, shape=(3, 2), transform=ZeroSumTransform(zerosum_axes))
        return m.logp()


# Works: zero-sum on the trailing axis
build_logp([1])
print("zerosum_axes=[1] (trailing): logp built OK")

# Fails: zero-sum on a leading axis -> ValueError: Incompatible Elemwise input shapes [(3,), (2,)]
build_logp([0])
print("zerosum_axes=[0] (leading): logp built OK")  # not reached

Error message:

Traceback (most recent call last):
  File "/tmp/zerosum_transform_bug.py", line 66, in <module>
    build_logp([0])
  File "/tmp/zerosum_transform_bug.py", line 58, in build_logp
    return m.logp()
           ^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pymc/model/core.py", line 728, in logp
    rv_logps = transformed_conditional_logp(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pymc/logprob/basic.py", line 642, in transformed_conditional_logp
    temp_logp_terms = conditional_logp(
                      ^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pymc/logprob/basic.py", line 572, in conditional_logp
    node_logprobs = _logprob(
                    ^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/functools.py", line 912, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pymc/logprob/transform_value.py", line 133, in transformed_value_logprob
    logprobs_jac.append(logp + log_jac_det)
                        ~~~~~^~~~~~~~~~~~~
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pytensor/tensor/variable.py", line 108, in __add__
    return pt.math.add(self, other)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pytensor/graph/op.py", line 209, in __call__
    node = self.make_node(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pytensor/tensor/elemwise.py", line 480, in make_node
    out_dtypes, out_shapes, inputs = self.get_output_info(*inputs)
                                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/velochy/miniconda3/envs/salk/lib/python3.12/site-packages/pytensor/tensor/elemwise.py", line 443, in get_output_info
    raise ValueError(
ValueError: Incompatible Elemwise input shapes [(3,), (2,)]

PyMC version information:

pymc 6.0.1

Context for the issue:

Currently bypassing this by not freezing my model before sampling, but it makes compile ~30% slower and likely slows down sampling as well.

Creating a PR with a proposed fix shortly.

Start in pymc/logprob/transform_value.py at transformed_value_logprob and inspect ZeroSumTransform.log_jac_det for the axis mismatch described. Reproduce the static shape case with ZeroSumTransform(zerosum_axes=[0]) and the freeze_dims_and_data path. Done means leading-axis zero-sum models build successfully without incompatible shapes, while trailing-axis behavior remains intact.