memory-efficient-attention (pypi): Memory Efficient Attention (O(sqrt(n)) for Jax and PyTorch
Keywords: attention|pytorch|jax