Skip to content

Add cross-attention to MultiheadAttention - #156

Merged
jessegrabowski merged 1 commit into
pymc-devs:mainfrom
jessegrabowski:cross-attention
Sep 5, 2026
Merged

jessegrabowski merged 1 commit into
pymc-devs:mainfrom
jessegrabowski:cross-attention

Conversation

@jessegrabowski

@jessegrabowski jessegrabowski commented Sep 5, 2026 •

Copy link
Copy Markdown
Member

MultiheadAttention.__call__ now takes an optional second input for the keys and values, with kv_dim for when it's a different width than the queries.

Self-attention output is bit-identical to before, checked against the pre-change code on fixed weights.


📚 Documentation preview 📚: https://pytensor-ml--156.org.readthedocs.build/en/156/

@jessegrabowski
jessegrabowski merged commit 92717e9 into pymc-devs:main Sep 5, 2026
11 checks passed
@jessegrabowski
jessegrabowski deleted the cross-attention branch September 5, 2026 06:07
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.

1 participant