Axial Positional Embedding

JAX implementation of lucidrains/axial-positional-embedding.

Installation

pip install axial-positional-embedding-jax

Citations

@inproceedings{kitaev2020reformer,
    title       = {Reformer: The Efficient Transformer},
    author      = {Nikita Kitaev and Lukasz Kaiser and Anselm Levskaya},
    booktitle   = {International Conference on Learning Representations},
    year        = {2020},
    url         = {https://openreview.net/forum?id=rkgNKkHtvB}
}

@misc{ho2019axial,
    title   = {Axial Attention in Multidimensional Transformers},
    author  = {Jonathan Ho and Nal Kalchbrenner and Dirk Weissenborn and Tim Salimans},
    year    = {2019},
    archivePrefix = {arXiv}
}