einexpr is a Python framework I wrote for working with named tensor dimensions. You can write the dimension names directly on arrays and use them to align, contract, rename and group dimensions.
import einexpr as ei
a = ei.array([1, 2])
b = ei.array([3, 4])
X = ei.array([[1, 2], [3, 4]])
Y = ei.array([[5, 6], [7, 8]])
dot = (a['i'] * b['i'])['']
outer = (a['i'] * b['j'])['i j']
product = (X['i j'] * Y['j k'])['i k']
The names determine how dimensions line up and which ones are contracted. In the matrix product, j appears in both inputs and disappears from the result. The expression does not depend on j occupying a particular axis.
Here is multi-head attention written both ways. PyTorch uses reshapes, transposes and axis numbers. einexpr names the dimensions on each array instead. t->s renames the key/value position before the contraction, and the result names say which dimensions remain.
The package also supported grouped dimensions. For example, context['b t (h d)'] joins the head and per-head feature dimensions while keeping the grouping explicit in the expression.