Skip to content

add protocol for array + offset -- to be used by RoPE utils#322

Merged
davidkoski merged 2 commits intomainfrom
positional-encoding
Jan 8, 2026
Merged

add protocol for array + offset -- to be used by RoPE utils#322
davidkoski merged 2 commits intomainfrom
positional-encoding

Conversation

@davidkoski
Copy link
Collaborator

Proposed changes

Adds a protocol:

public protocol OffsetLayer: Module {
    func callAsFunction(_ x: MLXArray, offset: Int) -> MLXArray
}

This will be useful in the RoPE utils in mlx-swift-lm.

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)

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.

👍

Comment on lines +1009 to +1012
/// A `Layer` (``Module`` subclass) that can be evaluated with an array and offset.
public protocol OffsetLayer: Module {
func callAsFunction(_ x: MLXArray, offset: Int) -> MLXArray
}
Copy link
Member

Choose a reason for hiding this comment

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

This is so specific to RoPE.. it might make sense to put it in the same file..

@davidkoski davidkoski merged commit b2cb6f6 into main Jan 8, 2026
7 checks passed
@davidkoski davidkoski deleted the positional-encoding branch January 8, 2026 21:13
@subsriram
Copy link

OffsetLayer is also declared in Module.swift:

public protocol OffsetLayer: Module {

@davidkoski this causes invalid redeclaration error

@davidkoski davidkoski mentioned this pull request Jan 8, 2026
4 tasks
davidkoski added a commit that referenced this pull request Jan 8, 2026
- build error as OffsetLayer was added in both #322 and #319
- #319 had it there "temporarily" for development but it should have been removed
- and #322 moved it to a new location so now we have two
- each PR built but the merge does not
@davidkoski davidkoski mentioned this pull request Jan 8, 2026
4 tasks
davidkoski added a commit that referenced this pull request Jan 9, 2026
- build error as OffsetLayer was added in both #322 and #319
- #319 had it there "temporarily" for development but it should have been removed
- and #322 moved it to a new location so now we have two
- each PR built but the merge does not
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.

3 participants