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