22 Jun, 2024 AcknowledgementsFirst, I would like to acknowledge my friends and kind internet strangers who helped me with this post. This post heavily adapts from the following - Tested Jax Transformer code via xjdr | Github Link Einstein summation in pytorch Einstein summation in numpy Basic guide to einsum Big thanks to _xjdr, Felix, Pushkar and Tokenbender for proof-reading. IntroI have been “delving” into jax and einsum notation lately in my quest to become a shape-rotator. This post is divi...