GitHub repository: ArvindSharma18/mesh-transformer-jax
Description: Model parallel transformers in JAX and Haiku