jax-vision-transformers Vision transformers with JAX & Flax Installation pip install jvt Usage example from jvt import DeiT ViT https://arxiv.org/abs/2010.11929 MAE https://arxiv.org/abs/2111.06377 DeiT https://arxiv.org/abs/2012.12877 LeViT https://arxiv.org/abs/2104.01136 ConvPass https://arxiv.org/abs/2207.07039