flax
feat(nnx): add preferred_element_type to attention module
#5253

Open

feat(nnx): add preferred_element_type to attention module #5253

aarushisingh04 wants to merge 4 commits into google:main from aarushisingh04:pr1-attention-preferred-element-type

feat(nnx): Add preferred_element_type to attention module

3d42df16

gemini-code-assist commented on 2026-02-15

Update flax/nnx/nn/attention.py

6672d14c

Update flax/nnx/nn/attention.py

f2de0876

Update tests/nnx/nn/attention_test.py

191250c7

Reviewers

gemini-code-assist

Assignees

No one assigned

Labels

None yet

Milestone

No milestone