Rewrite specifically for Sum and Prod to remove Join #951
Conversation
|
This may overlap/be redundant with #888 |
| @register_canonicalize | ||
| @register_uncanonicalize | ||
| @register_specialize | ||
| @node_rewriter([Sum, Prod]) |
There was a problem hiding this comment.
When there's nothing special about Sum and Prod we should apply the rewrites to all CAReduce operations, which also include stuff like Max, All, Any, ..., of which Sum/Prod are just two more instances
See the related PR I linked to in the comment
There was a problem hiding this comment.
Oops, 2nd time. I did not see this one show up in #59 😄.
There was a problem hiding this comment.
No they are not the same exactly.
There's a rewrite for reduction along axis0 for a join along axis 0.
My PR extends this to any axis.
There's then the question of multiple axis, of which axis=None is the most extreme (all axes). This PR can cover that case.
We may also want to think about multiple but not all axis. In what cases can we reduce first and join later?
Description
Rewrite specifically for Sum and Prod to remove Join. So graph that starts as
Sum{axes=None} [id A] 1 └─ Join [id B] 0 ├─ 0 [id C] ├─ x [id D] └─ y [id E]will become
Add [id A] 2 ├─ Sum{axes=None} [id B] 1 │ └─ x [id C] └─ Sum{axes=None} [id D] 0 └─ y [id E]Related Issue
Sums ofMakeVectors andJoins #59Checklist
Type of change