linalg:matmul

(linalg:matmul a b)

The matrix product of a and b (numpy's np.matmul, the @ operator). At rank <= 2 it behaves like linalg:dot -- matrix . vector included -- but signals an error when either operand is a scalar, catching the mistake of writing a matrix product where an elementwise linalg:mul would silently apply. The inner dimensions must agree; a mismatch signals an error.

At rank >= 3 on either side it is the stacked matrix product (torch's bmm / matmul): the last two axes are the matrix and every leading axis broadcasts by the numpy rules, so a (batch heads n d) query times a (batch heads d n) key gives (batch heads n n) attention scores. A rank-1 operand is promoted for the product -- a row on the left, a column on the right -- and its axis is dropped again from the result, exactly as in numpy.