Flash Attention Jax

https://github.com/lucidrains/flash-attention-jax

Flash Attention Jax is an implementation of the Flash Attention algorithm in Jax, providing...

Tokens:3,706
Snippets:32
Trust Score:9.9
License:MIT
Update:1 month ago
Tokens:
Raw