salesforce/jaxformer: Minimal library to train LLMs on TPU in JAX with pjit().
Category: llm
Language: Python