Facilitate graph pattern matching #1786
Draft
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR provides a POC for #1455
owner_opandowner_op_and_inputs.if x.owner.op is not None and isinstance(x.owner.op, ...)in favor ofisinstance(x.owner_op, ...)casestatement. This makes nested patterns easier to write. Eg.case (Blocwise(Dot()), a, b)for a singlematch x.owner_op_and_inputs.match_argsso common parametrized Ops (Elemwise, Blockwise) and Constants are easier to pattern match without kwargsis_matrix_transposeandis_left_expanded_matrix_transposeproperty to DimShuffle, for easier pattern matchingI then refactored the existing linalg rewrites as a POC. I went a bit overboard and tried to standardize names, and fuse rewrites that were pretty simple to combine.
CC @tanish1729 who did many of these rewrites, see if it feels nicer?