Nice! I remember trying to implement Hinton's capsules layer in tensorflow [1], and stumbling upon Einstein notation syntax. Took me a while to figure it out.
[1] https://github.com/michaelklachko/CapsNet/blob/master/capsne...
[1] https://github.com/michaelklachko/CapsNet/blob/master/capsne...
The same convention as in NumPy, but IMHO much nicer visual explanation for beginners. Also, think about tensor diagrams, as in the Feynman diagram in "Simple diagrams of convolutional neural networks" https://medium.com/inbrowserai/simple-diagrams-of-convoluted....