Einsum notation gives each input axis a letter, then expresses multiplication and summation in one compact equation. It makes complex array operations easier to read, especially inside Transformer attention code. section 15.3
Read an equation using these rules:
- Write one label for each axis of every input, such as ijk for a rank-3 input. section 15.3
- Multiply values along repeated labels across inputs; the corresponding dimensions must match. section 15.3
- Sum over labels that appear in the inputs but not in the output. section 15.3
- Arrange returned axes in the order written after the arrow. section 15.3
Writing dot-product attention
For target and source sequences, scores = np.einsum("btd,bsd->bts", target, source) multiplies matching feature dimensions and produces a score for every target–source pair. After softmax, np.einsum("bts,bsd->btd", scores, source) forms the weighted source representation for each target position. section 15.3.1
The same notation supports scalar sums, transposition, matrix multiplication, and higher-rank operations. In Keras, keras.ops.einsum replaces np.einsum, while EinsumDense uses einsum instead of the matrix multiplication in a Dense layer. section 15.3
np.einsum("ij->ji") #1 np.einsum("ij,jk->ik") #2 np.einsum("hij,jk->hik") #3 np.einsum("i,i->") #4 np.einsum("ijk,ijk->ijk") #5 np.einsum("ijk,ijk->") #6
Use the compact listing to compare transposition, matrix multiplication, scalar reduction, elementwise operations, and full reductions. section 15.3
