Broadcasting by name
The order rule
Elementwise operations (+ - * /, comparisons, and, or, and elementwise built-ins and functions) line their operands’ axes up as follows. A rank-0 operand broadcasts to anything. Otherwise:
- By name. If every axis of one operand has a name and those names are all axes of the other operand, the axes line up by name. The result has the other operand’s axes, in its order. If both operands qualify, as with
[n m] + [m n], the left operand’s order wins. The matched axes must have equal sizes (E0401). If a name appears several times, theith axis with that name in one matches theith in the other. - From the right. Otherwise the axes line up from the right, as in NumPy. Each pair must have the same size, or one of them must have size 1, which broadcasts (E0401). Two axes in a pair that both have names must have the same name (E0404). The result takes the size of each pair (a named axis is kept over an unnamed one of the same size), plus the extra leading axes of the longer operand.
So [b H n m] + [n m] lines up by name and gives [b H n m], as does [n m] + [b H n m]. [n D] + [k D] is E0404 even when n and k have the same size: rename one with rearrange or as if they really are the same axis. [n D] + [D] lines up by name; [n 3] + [3] lines up from the right.
The generated Python permutes and unsqueezes operands so they line up as the checker says, so no unsqueeze appears in the source. masked_fill requires its mask to broadcast to the tensor’s shape without adding axes (E0404).
Matrix multiplication
x @ y and matmul(x, y) work like torch.matmul. The last axis of x is contracted with the second-to-last axis of y, or its only axis if it has one. The contracted axes need equal sizes but not equal names (E0401), so [b n (H Dh)] @ [D D] is fine when H*Dh = D. With two or more axes each, the axes before the last two broadcast by the order rule, and the result is those, then x’s second-to-last axis, then y’s last. A rank-1 operand loses its contracted axis, and a rank-0 operand is E0403.
Both operands must have the same dtype (E0502), since torch doesn’t promote for @; the result is Real.
Operators
| Operator | Operands | Result |
|---|---|---|
+ - * / | numbers, not Mask (E0505) | numbers; range from range inference; / of integers gives the default float |
@ | numbers of one dtype | Real |
== != < <= > >= | anything | Mask; a bool for two host scalars |
and, or | Masks, or bools (E0501) | Mask, or bool |
! | a Mask or a bool (E0501) | Mask, or bool |
unary - | numbers (E0505) | numbers; negated logits stay logits |
An operand that is a Python value makes the whole result a Python value (Python interop). A string operand is E0412.