Avoid casting all terms to the same dtype in logp

[ricardoV94]

Description

This change avoids many warnings when compiling a logp to JAX in pytensor.config.floatX = "float32". Since the terms are all scalars, it should also be more efficient, because it avoids an explict MakeVector.

Related Issue

• Closes #
• Related to #

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):

---

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

[lucianopaz]

Why does this change affect int casting? Aren’t the logp terms already floatX?

[ricardoV94]

Why does this change affect int casting? Aren’t the logp terms already floatX?

This change avoids explicit castings to float64. Some logps may be float32 and others float64, and if you sum there will be an explicit casting + make vector. If you do add there is only implicit casting and jax doesn't emit a warning about ignoring explicit float64 casts.

This has nothing to do with ints anymore, I have split those changes into another PR

[lucianopaz]

LGTM. The errors seem like a pytensor problem with ScalarType.upcast, so this PR should not be merged until pytensor sorts that out.

[ricardoV94]

Not a PyTensor error, but a cryptic failure when add has no entries

[ricardoV94]

Actually we have a rewrite that does this, but it's a bit conservative with dtypes... Maybe we can tweak this rewrite instead.

https://github.com/pymc-devs/pytensor/blob/f7b0a7a48b929605a083e13a12f144040a7fe265/pytensor/tensor/rewriting/basic.py#L919-L956