A collection of model architectures implemented in jax/flax with training code. Convolutional Autoencoder with Attention Vision Transformer Sources List of useful resources http://jalammar.github.io/illustrated-transformer/ https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial2/Introduction_to_JAX.html