flash-attention-jax (pypi): Flash Attention - in Jax
Keywords: artificial intelligence|deep learning|transformers|attention mechanism|jax|artificial-intelligence|attention-mechanisms|deep-learning|long-context-attention