GitHub repository: vyomkeshj/mesh-transformer
Primary language: Python
Description: Model parallel transformers in JAX and Haiku