blog

my blog about machine learning and life


JAX-LM: Language Modelling and Distributed Training in JAX

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

78 min read   ·   2026