flax

e95acab1 - Merge pull request #1291 from avital:attention-weights