Skip to content

Conversation

@ricardoV94
Copy link
Member

@ricardoV94 ricardoV94 commented Dec 12, 2025

This PR provides a POC for #1455

  • It adds two variable properties: owner_op and owner_op_and_inputs.
    • This allows us to avoid the verbose if x.owner.op is not None and isinstance(x.owner.op, ...) in favor of isinstance(x.owner_op, ...)
    • The second one is more niche, but allows to both match an Op and structurally unpack the inputs in case of a match in a single case statement. This makes nested patterns easier to write. Eg. case (Blocwise(Dot()), a, b) for a single match x.owner_op_and_inputs.
  • It adds match_args so common parametrized Ops (Elemwise, Blockwise) and Constants are easier to pattern match without kwargs
  • It adds a much needed is_matrix_transpose and is_left_expanded_matrix_transpose property to DimShuffle, for easier pattern matching

I 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?

# Check for use of pt.diag first
case (AllocDiag(offset=0, axis1=axis1, axis2=axis2), diag_input):
ndim = diag_input.ndim
if axis1 == ndim - 1 and axis2 == ndim:
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is new, we were not checking for non standard allocdiagonal, not along the last dims!

@ricardoV94 ricardoV94 force-pushed the op_shortcut branch 2 times, most recently from 0a5ec73 to cf91364 Compare December 12, 2025 14:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant