An open-source implementation of the Rethinking Spatial Dimensions of Vision Transformers research paper in Google's #JAX and #Flax.

JAX/Flax Repository: https://github.com/conceptofmind/PiT-flax

PiT Research Paper: https://arxiv.org/abs/2103.16302

Official Github Repository: https://github.com/naver-ai/pit

In collaboration with Dr. Phil 'Lucid' Wang: https://github.com/lucidrains.