RetNet in JAX

Paper: Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., ... & Wei, F. (2023). Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621. Inspiration: https://github.com/Jamie-Stirling/RetNet/tree/maiN

• NOTE: This is meant as a personal learning project, so implementation may not be optimal and no guarantees on correctness of code...

Usage

CPU:

$ pip install -r requirements.txt

GPU:

$ pip install -r requirements_gpu.txt -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html