flax

e0fa96fe - Merge pull request #3835 from google:linen-attention-multiple-initializers