Fix per-example mask + docs in sdpa

ml-explore / mlx

MLX: An array framework for Apple silicon

https://ml-explore.github.io/mlx/

MIT License

17.37k stars 1.01k forks source link

Closed awni closed 1 week ago

awni commented 2 weeks ago

This worked in the regular case but was broken in the gqa case since we reshape to 5d.
FIx + test
More in the docs on this.