Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rewrite scalar dot as multiplication #1205 #1263

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

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