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