Skip to content

Conversation

SwarnimShekhar
Copy link

@SwarnimShekhar SwarnimShekhar commented Mar 4, 2025

Description

  • Optimized local_blockwise_dot_to_mul to handle (1,1) @ (1,1) cases as direct elementwise multiplication.
  • Prevents unnecessary reshaping for scalar-like tensors.
  • Ensures consistency with other elementwise multiplication cases.

Related Issue

  • Closes #
  • Related to #

Checklist

Type of change

  • New feature / enhancement
  • Bug fix
  • Documentation
  • Maintenance
  • Other (please specify):

📚 Documentation preview 📚: https://pytensor--1263.org.readthedocs.build/en/1263/

Comment on lines +299 to +301
# new condition to handle (1,1) @ (1,1)
if a.ndim == 2 and b.ndim == 2 and a.shape == (1, 1) and b.shape == (1, 1):
return [a * b] # Direct elementwise multiplication
Copy link
Member

Choose a reason for hiding this comment

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

This was already handled, the issue that was not handled was for the care dot (not the blockwised versions) that this rewrite tracks.

Also a.ndim is redundant if you are checking the shape explicitly. More important you have to use static shape a.type.shape not the symbolic shape a.shape, as is done in the rest of this function

Copy link
Author

Choose a reason for hiding this comment

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

I will look into that again!

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.

2 participants