BUG: ZeroSumTransform jacobian shape breaks on model freeze

[velochy]

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.

[ricardoV94]

you cannot zerosum on the leading axes, core dims must be rightmost

[ricardoV94]

you're also not supposed to use this with scalar distributions, it's a vector/ndim transform, really only meant for zerosumnormal, but it should work shapewise if you zerosum on the right

[velochy]

So why does it have the misleading zerosum_axes argument that allows me to sum on non-leading axes?
And I do have use for it and it's served me pretty well so far :)

Also, for the record, this only breaks when freezing the model. with non-frozen model, things work as they are :)

[velochy]

Anyway: your call. If you decide not to support those use cases I will roll my own local version and move on. I already have an extended-functionality MultiZeroSum so it would basically just require changing some references.

[ricardoV94]

So why does it have the misleading zerosum_axes argument that allows me to sum on non-leading axes?

A design mistake imo. We at least forced this away in the ZeroSumNormal itself, with n_zerosum_axes. I think initially it had arbitrary axes and so did the transform, which stayed in the original form.

And I do have use for it and it's served me pretty well so far :)

Yeah it may be fine to use, I meant more to say it was not designed for that, but you're free to use transforms that distort other distributions, not just in those where it keeps the same as prior. Just the usual remark that prior_predictive won't respect it.

this only breaks when freezing the model. with non-frozen model

You are seeing pytensor optimizing under the assumption you have a valid shape. Usually you can "see" the error if you compile with get_mode().excluding("shape_unsafe"), although we are not 100% thorough on marking rewrites that make such shape assumptions, but we mark the most obvious ones. (You never want to run with that, it's just for debug if you have a shape error and pytensor is sweeping it under the rug).

In your cases you have add(x, zeros(n, m)), there's no broadcasting of x with zeros, so it just drops them.

If you decide not to support those use cases I will roll my own local version and move on.

We do rely on the fact that core dimensions are rightmost, keeps code simpler throughout, and actually correct for arbitrary logprob inference, as we would otherwise need #6360. This is documented on the forward side here, but same applies to logp, when dimensions drop for regular distributions, they are always on the right: https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/dimensionality.html#multivariate-distribution-example

To be on the safe side, I suggest you just transpose your variable so the core_dims/zerosum_axes are on the right, then transpose back afterwards for whatever use you have of it. Same thing if you wanted a Dirichlet with probabilities on any axis besides the last...

[ricardoV94]

What does MultiZeroSum do?

[velochy]

What does MultiZeroSum do?

Are you sure you want to know?

[velochy]

It basically allows defining multiple ranges over the last axis that each have a zero-sum constraint.
And most recently, I added a feature where those ranges can be flattened n-dimensional tensors where I need zerosum over multiple axes of the tensor.

As for why I would need something like this: you yourself told me to vectorize my models, so I did :D
(I'm basically stacking multiple models on top of each other, so if I have multiple multinomials I'm observing, the effects matrices need multiple zerosum ranges. As for the tensor extension - trust me, you don't want to know)