Nair, L. (2024). Can FlashAttention Statistics be Learned? (Version 2.0.4) [Computer software]. https://github.com/lnairGT/Learned-Flash-Attention