[P] deit3-jax: A codebase for training ViTs on TPUs

Hey all, I have written a codebase to train ViTs by following DeiT and DeiT-III recipes. As they are strong baselines to train vanilla ViTs, it is necessary to reproduce to adopt to the variant research. However, the original repository is implemented in PyTorch, it is impossible to run on TPUs.

Therefore I re-implemented the simple ViT training codebase with DeiT and DeiT-III training recipes. Here is my repository: [https://github.com/affjljoo3581/deit3-jax](https://github.com/affjljoo3581/deit3-jax). I used Jax/Flax and webdataset to build a TPU-friendly training …

codebase hey however iii jax machinelearning pytorch recipes research tpus train training