GitHub repository: 47-studio-org/mesh-transformer-jax-2
Description: Model parallel transformers in JAX and Haiku