Prevent local_sum_make_vector from introducing forbidden float64 at graph level

[tvwenger]

Description

This PR changes the behavior of local_sum_make_vector to skip rewriting whenever the internal accumulator is more precise than the input/output data. Otherwise, the internal accumulator is added to the graph, which can introduce forbidden float64 to the graph when the user requests config.floatX="float32" or less precision.

Replaces PR #655

Related Issue

• Closes local_sum_make_vector rewrite can introduce forbidden float64 operations at the graph level  #653

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

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 80.82%. Comparing base (e8693bd) to head (d68a252).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #656   +/-   ##
=======================================
  Coverage   80.82%   80.82%           
=======================================
  Files         162      162           
  Lines       46812    46814    +2     
  Branches    11437    11438    +1     
=======================================
+ Hits        37836    37838    +2     
  Misses       6725     6725           
  Partials     2251     2251           

Files	Coverage Δ
pytensor/tensor/rewriting/basic.py	94.07% <100.00%> (+0.02%)	⬆️

[tvwenger]

@ricardoV94 This PR is ready for review.

[ricardoV94]

This seems less restrictive and still avoid the original issue?

Suggested change

	if is_an_upcast(out_dtype, acc_dtype):
	if acc_dtype=="float64" and out_dtype != "float64" and pytensor.config.floatX != "float64":

[tvwenger]

The issue could also appear if the user requests even less precision. For example, if config.floatX = "float16", then acc_dtype will be float32 and we'll have the same problem. The issue will always exist when acc_dtype is more precise than out_dtype, which is true for every floatX except float64.

[tvwenger]

I take your point though. I'll re-work this.

[ricardoV94]

Do we support floatX=16 btw? I think float16 was only ever an experimental feature in the codebase and not widely supported

[tvwenger]

@ricardoV94 This will be ready for review when the checks are complete

[ricardoV94]

This is still too restrictive. floatx=32 does not mean float64 are forbidden, only that PyTensor won't introduce them automatically if they can be avoided. The user is still allowed to call sum(dtype=float64) or pass a float64 input in which case this check is too restrictive. There's already a float64 at one of the ends of this operation so it's fine for the float64 acc_dtype to be exposed.

That's why my suggestion was quite more verbose.