Port Vision Transformer (ViT) to Objax examples

[david-berthelot]

https://github.com/google-research/vision_transformer

[sshkhr]

@david-berthelot Is there still interest in porting ViT to Objax examples? I am new to JAX and I'd be interested in adding Vit to the objax zoo (and examples)

[AlexeyKurakin]

@sshkhr
It would be great if you can volunteer to do this.
As you mentioned, the model itself could be added to objax zoo, and training code to examples

[sshkhr]

Sounds good @AlexeyKurakin! It will take me a few days to get started on this but I will send in a PR early next month with the model.