flax

a59675c1 - Merge pull request #3566 from google:nnx-attention-add-qkv-features