Ported weights from official JAX impl of Vision Transformers and MLP-Mixer

Converted to PyTorch from https://github.com/google-research/vision_transformer