Build a Transformer in JAX from scratch: how to write and train your own models

How to develop and train a Transformer with JAX, Haiku and Optax. Learn by example how to code Deep Learning models in JAX