JAX implementation of ViT-VQGAN.
- Jiahui Yu for his advice on the implementation based on the work on ViT-VQGAN and Parti.
- Phil Wang for suggesting using convolutions after self-attention layers.
- Katherine Crowson for suggesting using convolutions in feed-forward layers.