Vanilla Transformer jax

• this notebook shows the implementation of Vanilla Transformer.
  1. Implement Vanilla Transformer with JAX/Flax.
  2. Using the Multi30k dataset, train Transformer to translate German to English.
  3. Translate some sentences with a greedy search translator.
  4. Plot attention matrices.
• Detailed explanation Blog Post (ja): JAX/Flaxでゼロから作るTransformer① ―Vanilla Transformer
• Detailed explanation Blog Post (en): Transformers from scratch with JAX/Flax① ―Vanilla Transformer