Skip to content

fix #315 -- do not reshape during SDPA#316

Merged
davidkoski merged 1 commit intomainfrom
rope-noreshape
Dec 15, 2025
Merged

fix #315 -- do not reshape during SDPA#316
davidkoski merged 1 commit intomainfrom
rope-noreshape

Conversation

@davidkoski
Copy link
Collaborator

@davidkoski davidkoski commented Dec 15, 2025

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

- this was changed on the python side in ml-explore/mlx@16546c7
- in v0.29.1 this still produced the correct result but in v0.30.0 it does not
	- in particular a different specialized function is used
@davidkoski davidkoski requested a review from awni December 15, 2025 21:53
public func callAsFunction(_ x: MLXArray, offset: Int) -> MLXArray {
let shape = x.shape
var x = x.reshaped(-1, x.dim(-2), x.dim(-1))
x = MLXFast.RoPE(
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Adopt ml-explore/mlx@16546c7 (fro Mar 18, 2024!!!)

Copy link
Member

@awni awni left a comment

Choose a reason for hiding this comment

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

Great catch!

@davidkoski davidkoski merged commit 62d4b95 into main Dec 15, 2025
5 checks passed
@davidkoski davidkoski deleted the rope-noreshape branch December 15, 2025 22:35
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