Expose generalized projection params in attention_pair_bias wrapper - #293
Merged
Conversation
Expose the optional generalized (Proteina/Complexa) projection parameters of the backend attention_pair_bias op as keyword-only arguments on the cuet frontend wrapper: b_proj_k, b_proj_v, w_ln_q/b_ln_q, w_ln_k/b_ln_k, and b_proj_g. They default to None, so the strict OpenFold3/Boltz contract applied to single_repr is unchanged; supplying them adds the K/V/gate projection biases and the post-projection Q/K LayerNorm (over the full H*DH dimension, before the head split). - primitives/triangle.py: add the 7 keyword-only params, document them, and forward them to the backend. - tests/primitives/triangle_test.py: add a case asserting the generalized params are forwarded (output differs from the strict call). - CHANGELOG.md, SKILL.md: document the new optional params. Depends on the backend generalization (kernelcatcher !530) being merged and a new cuequivariance_ops_torch wheel published; public CI installs the released backend wheel, so this cannot go green until then. Torch-only; there is no JAX cuex.attention_pair_bias counterpart. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Alexandra Kim <alexandrak@nvidia.com>
jomitchellnv
approved these changes
Jul 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Extends the Torch
attention_pair_biaswrapper to expose the generalized backend API:w_proj_z.cuequivariance_ops_torchAPI instead of its internal module path.Torch only; there is no JAX counterpart. This requires a backend release with the generalized
attention_pair_biassignature.Validation
pytest cuequivariance_torch/tests/primitives/triangle_test.py -q— 7 passedpre-commit run --all-files— passedChecklist
pre-commitlocally.