Transformers from scratch with JAX/Flax

[kenkenpa2126]

Hi, I'm Kenta, a student in Japan.
I started a tutorial to understand Transformer models through implementation with JAX/Flax.
As the first step, even though I haven't finished writing detailed explanations yet, I released the notebook of Vanilla Transformer.
https://github.com/kenkenpa2126/Transformers_from_scratch_with_JAX-Flax

I aim to register it to the community example as a tutorial and wonder how to apply it.
Is it OK to list it on this page directly and make a PR?
https://github.com/google/flax/blob/main/docs/examples_community_examples.rst

And also, If you gave me some advice or comments about the tutorial, I would appreciate it very much!

[8bitmp3]

It looks like a good candidate for this section here: https://flax.readthedocs.io/en/latest/examples_community_examples.html#examples
@andsteing @marcvanzee

[chiamp]

Also, I believe @levskaya and @IvyZX were working on a Flax implementation of a transformer?

[IvyZX]

Yes, but our minimal transformer will be decoder-only while this one is encoder-decoder. It could be part of the community tutorial section too!

[kenkenpa2126]

Thanks for your comments! I'll make a PR after I write some explanation.