Preface This is a series of a tutorial to understand the implementation of Transformer Models with JAX/Flax. As a first step, this blog post shows how to implement Vanilla Transformer. Vanilla Transformer, aka Transformer Encoder-Decoder M…

Preface このシリーズは、JAX/FlaxでゼロからTransformer系のモデルを実装し、JAX/Flaxの使い方やモデルの仕組みを理解するためのチュートリアルである。この記事では、その第一弾として、最初のTransformerであるVanilla Transformerを実装する。Vanilla Tr…