JAX-LM: Language Modelling and Distributed Training in JAX

A guide to language modelling and distributed training from scratch in JAX