Metadata-Version: 2.2
Name: axial-positional-embedding-jax
Version: 0.0.1
Project-URL: Homepage, https://github.com/axrwl/axial-positional-embedding
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Requires-Python: >=3.8
Description-Content-Type: text/markdown

# Axial Positional Embedding

JAX implementation of [lucidrains/axial-positional-embedding](https://github.com/lucidrains/axial-positional-embedding).

### Installation
`pip install axial-positional-embedding-jax`

### Citations

```bibtex
@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}
}
```

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