Fix ZeroSumTransform logp for non-trailing zero-sum axes

[velochy]

Description

transformed_value_logprob reduces the logp's trailing axes to match a lower-dim Jacobian, assuming (as for Simplex/Ordered) that a multivariate transform acts on the trailing axes. ZeroSumTransform may zero-sum arbitrary axes, so for a non-trailing axis the reduced logp keeps the wrong axis and disagrees with log_jac_det (which drops zerosum_axes). The shapes then differ, e.g. (3,) + (2,). This is silent with symbolic dims (both broadcastable) but raises Incompatible Elemwise input shapes once dims are static -- e.g. after freeze_dims_and_data, which the numba/jax sampling backends apply.

Fix: expose the integrated-out axes via ZeroSumTransform.jacobian_reduce_axes and have transformed_value_logprob reduce those axes (falling back to the trailing axes for transforms that don't define them). ZeroSumNormal is unaffected (it only zero-sums trailing axes); batched transformed-logp shapes are preserved.

Minimal repro (raised before, builds after)::

    with pm.Model() as m:
        pm.Normal("x", 0.0, 1.0, shape=(3, 2),
                  transform=ZeroSumTransform([0]))
        m.logp()

Related Issue

• Closes BUG: ZeroSumTransform jacobian shape breaks on model freeze #8335

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[velochy]

@ricardoV94 not sure this is the solution you want (requires adding a new property to transform), but no better ones immediately come to mind. Feel free to propose a different approach if you see one.

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #33237924 | 📁 Comparing b27dff9 against latest (94787c0)

  🔍 Preview build  

3 files changed

± glossary.html
± _modules/pymc/distributions/transforms.html
± api/distributions/generated/pymc.distributions.transforms.ZeroSumTransform.html

[codecov]

Codecov Report

❌ Patch coverage is 85.71429% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 87.28%. Comparing base (9d24260) to head (b27dff9).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/distributions/transforms.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8336      +/-   ##
==========================================
+ Coverage   79.80%   87.28%   +7.48%     
==========================================
  Files         125      125              
  Lines       20526    20532       +6     
==========================================
+ Hits        16380    17922    +1542     
+ Misses       4146     2610    -1536     

Files with missing lines	Coverage Δ
pymc/logprob/transform_value.py	96.69% <100.00%> (+0.08%)	⬆️
pymc/distributions/transforms.py	99.53% <66.66%> (-0.47%)	⬇️

... and 41 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.