I notice you can use splash attention from jax-ml (flash) or tokamax (tokamax_flash). Im wondering why there's two different sources for it, and which would be recommended? Glancing at the two versions, the tokamax one seems to be a bit more up to date? But wanted to double check. Thanks.
