GitHub repository: Abaso007/jaxformer
Description: Minimal library to train LLMs on TPU in JAX with pjit().