最近、Jax/Flaxを触るようになりました。ここでは、Jax/Flaxを用いてVision Transformerを実装した方法と、Jax/Flaxによる学習の方法について紹介しようと思います。 Vision Transformerのおさらい Vision Transformerを実装するにあたって、まずはこの図を…