Cast ZeroSumNormal shape operations to config.floatX

[thomasjpfan]

What is this PR about?
This PR prevents ZeroSumNormal from upcasting to float64, when pytensor is configured with floatX = float32.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Bugfixes

• Fixes BUG: ZeroSumNormal unintentionally upcasts to float64 #6886

---

📚 Documentation preview 📚: https://pymc--6889.org.readthedocs.build/en/6889/

[codecov]

Codecov Report

Merging #6889 (2d1546d) into main (dfb05b6) will decrease coverage by 0.37%.
Report is 2 commits behind head on main.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6889      +/-   ##
==========================================
- Coverage   92.05%   91.69%   -0.37%     
==========================================
  Files          96      100       +4     
  Lines       16446    16851     +405     
==========================================
+ Hits        15140    15451     +311     
- Misses       1306     1400      +94     

Files Changed	Coverage Δ
pymc/distributions/multivariate.py	92.55% <100.00%> (ø)
pymc/distributions/transforms.py	99.37% <100.00%> (+<0.01%)	⬆️

... and 12 files with indirect coverage changes

[ricardoV94]

Thanks, just a small suggestion to make the test faster

[ricardoV94]

Should be enough to call model.logp()

[thomasjpfan]

I do not think it is enough because logp does not call ZeroSumTransform.backward. pm.sample is a way to test both ZeroSumTransform.backward and logp. I added a comment about this interaction in 7906292 (#6889)

Note that the new test is mostly a non-regression test for #6886.

[thomasjpfan]

Ah I see how m.logp works. Thanks for the suggestion.

[ricardoV94]

I suggested just calling model.logp because it's enough to identify the original problem and it's much faster than compiling and sampling a model. This test is too heavy for the change being made.

[thomasjpfan]

Okay, I updated the test to use logp.

[ricardoV94]

Sorry I wasn't clear. I meant you should call model.logp().

import pymc as pm
import pytensor 

with pytensor.config.change_flags(floatX="float32", warn_float64="raise"):
    with pm.Model() as m:
        pm.ZeroSumNormal("b", sigma=1, shape=(2,))
    m.logp()

[ricardoV94]

That should call both backward and forward

[ricardoV94]

Thanks !!!

[ricardoV94]

Thanks @thomasjpfan!