soxueren/vit-jax-flax

Vision Transformer from scratch (JAX/Flax).

vit-jax-flax

Vision Transformer from scratch (JAX/Flax).

Implementation details are described in Vision Transformer from scratch (JAX / Flax)