Minimal but scalable implementation of large language models in JAX - View it on GitHub