Skip to content

More primitives for compiling with shapeless #1653

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.

Sign up for GitHub

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

Jump to bottom
Merged
merged 4 commits into from
Dec 6, 2024
Merged

More primitives for compiling with shapeless #1653

merged 4 commits into from
Dec 6, 2024

Conversation

awni
Copy link
Member

@awni awni commented Dec 6, 2024 •
edited
Loading

Added several more primitives to allow compiling with shapeless=True (and hopefully exporting with shapeless=True).

  • Arange
  • Gather
  • Transpose
  • Concatenate
  • Matmul
  • QuantizedMatmul
  • fast::AffineQuantize
  • fast::LayerNorm
  • fast::RMSNorm
  • fast::RoPE
  • fast::ScaledDotProductAttention

@awni awni requested review from angeloskath and barronalex and removed request for angeloskath December 6, 2024 19:01
angeloskath
angeloskath approved these changes Dec 6, 2024
View reviewed changes
Copy link
Member

@angeloskath angeloskath left a comment

Choose a reason for hiding this comment

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

Looks good. There isn't really a reason to not do a shapeless compile for all primitives is there? Except having to write output shape functions.

@awni
Copy link
Member Author

awni commented Dec 6, 2024

There isn't really a reason to not do a shapeless compile for all primitives is there?

I think for almost all primitives we should allow it.

A couple primitives are somewhat dangerous to allow for shapeless compile. Mostly Reshape and Broadcast where the output shape is hardcoded and there is not currently a way to infer it from the input shape. It might still be nice to allow them to be shapelessly compiled because you may know that the shapes for that part of the graph are fixed.. something to think about.

@awni awni merged commit d0b6cb0 into main Dec 6, 2024
5 checks passed
@awni awni deleted the more_shapeless branch December 6, 2024 19:29
@BrewTestBot BrewTestBot mentioned this pull request Jan 9, 2025
Merged
1 task
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@angeloskath angeloskath angeloskath approved these changes

@barronalex barronalex Awaiting requested review from barronalex

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@awni @angeloskath