Merge consecutive reduces

Open ricardoV94 opened this issue 1 year ago • 1 comments

Description

This PR does some cleanup and extends reduction rewrites. Check each commit for the individual changes. Summary:

Remove list of ALL_REDUCE. AFAICT the tracks machinery already handles subclasses automatically. It did not always do it, which is probably why the list existed
~~Stop defining Min as negative of Max, to later undo it and actually use Min. I think this is just some historical accident?~~ Reverted, too many things depend on this for me to tackle now: #904
Generalize a chain of reductions of the same type to any CAReduce. It was previously limited only to Sum and Prod, but I can't see a good reason for it. This showed up in vectorized logps where All reductions (related to the parameter checks) were not being combined.
Generalize local_reduce_join to any axis and regardless of whether the inputs along that axis are broadcastable because of expand_dims or any other reason. It shouldn't matter, as long as we can squeeze it out.

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

Jul 05 '24 19:07 ricardoV94