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