Skip to content

Add RoPE embeddings#5481

Merged
copybara-service[bot] merged 1 commit into
google:mainfrom
samanklesaria:rope
Jul 23, 2026
Merged

Add RoPE embeddings#5481
copybara-service[bot] merged 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!

@copybara-service
copybara-service Bot merged commit a525264 into google:main Jul 23, 2026
19 of 26 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants