SliceOut: Training Transformers and CNNs faster while using less memory. - Pascal Notin; Aidan N. Gomez; Joanna Yoo; Yarin Gal (CoRR, 2020)

Scalable Training of Language Models using JAX pjit and TPUv4. - Joanna Yoo; Kuba Perlin; Siddhartha Rao Kamalakara; João G. M. Araújo (CoRR, 2022)