Simple Transformer in Jax
This is a teaching tool. It has been tested to run correctly. It is not complete. Thair still be dragons here. Please act accordringly ... check out https://sankalp.bearblog.dev/einsum-new/ for a more detailed explanation fow how and why this works.