Vision Transformers in JAX

Transformers were introduced in 2017 and changed the course of Deep Learning research for many forthcoming years. Naturally, they were also adopted in the computer vision ecosystem which was till then, dominated by Convolutional Neural Networks or Co...