Module Nx.Pattern
Axis patterns.
type tThe type for axis patterns, parsed and checked.
v s is the pattern s. Each call parses s; a pattern bound at a module's top level is parsed once, as the module starts.
With one operand that repeats no name and sums none, every name appears once on each side; a group on the left splits an axis and on the right merges axes; 1 drops a unit axis on the left and adds one on the right; ... is the same axes on both sides.
Otherwise, a name in the result is a batch axis where two operands or more hold it and a free axis where one does, and the names after | are summed: each is in an operand and not in the result. Every operand name is in the result or summed. A name an operand repeats takes one index on all of its axes there, so "b i i -> b i" is the diagonal of each matrix. A group in an operand splits that operand's axis, and an extent a group leaves unknown comes from an axis the name stands on alone, in any operand, or from sizes. ... stands for the same leading axes in every layout, kept as batch axes; 1 drops a unit axis of an operand or adds one to the result.
Sums are written out because a name dropped from the result is easy to miss. In "b h q d, b h k d -> h q k | d", b was left out by mistake; summing it would give a result of plausible shape and wrong values, so the pattern raises.
Raises Invalid_argument naming s if it is not in the grammar, a name repeats in the result or after |, ... appears twice in a layout, or a name breaks the rules above. A string in NumPy's einsum notation, such as "bik,bkj->bij", raises with the pattern it means: "b i k, b k j -> b i j | k".