Skip to content

Add RoPE embeddings#5481

Open
samanklesaria wants to merge 1 commit into
google:mainfrom
samanklesaria:rope
Open

Add RoPE embeddings#5481
samanklesaria wants to merge 1 commit into
google:mainfrom
samanklesaria:rope

Conversation

@samanklesaria

Copy link
Copy Markdown
Collaborator

This PR adds support for RoPE. Specifically, a new function dot_product_attention_with_rope can be used as the attention_fn argument for nnx.MultiHeadAttention.

Comment thread flax/nnx/nn/attention.py
Comment thread flax/nnx/nn/attention.py Outdated
@samanklesaria samanklesaria force-pushed the rope branch 3 times, most recently from 319550d to ea04909 Compare June 3, 2026 18:04
@samanklesaria samanklesaria requested a review from vfdev-5 June 3, 2026 18:06
Comment thread flax/nnx/nn/attention.py Outdated
Comment thread flax/nnx/nn/attention.py
@samanklesaria samanklesaria force-pushed the rope branch 2 times, most recently from c7b1f41 to efa289d Compare June 16, 2026 18:02
Comment thread tests/nnx/nn/attention_test.py
Comment thread tests/nnx/nn/attention_test.py
@vfdev-5

vfdev-5 commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator

@samanklesaria can you fix the merging conflict, thanks!

@vfdev-5

vfdev-5 commented Jun 30, 2026

Copy link
Copy Markdown
Collaborator

@samanklesaria can you check mypy CI job if it is real? Thanks!

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