flax

a5661961 - Merge pull request #3687 from chiamp:nnx_attention