Matrix × matrix

$A$ is $2 \times 3$ and $B$ is $3 \times 2$. $AB$ will have $A$’s 2 rows and $B$’s 2 columns. Each of its 4 entries is one dot product.

Row 1 of $A$ dotted with column 1 of $B$: $1 \cdot 2 + 2 \cdot 1 + 0 \cdot 4 = 4$.

Same row, next column: $1 \cdot 1 + 2 \cdot 0 + 0 \cdot (-1) = 1$.

Row 2 with column 1: $3 \cdot 2 + (-1) \cdot 1 + 1 \cdot 4 = 9$. Row 2 with column 2: $3 \cdot 1 + (-1) \cdot 0 + 1 \cdot (-1) = 2$.

The inner sizes must match, because each row of $A$ is dotted with each column of $B$. They disappear in the result; the outer sizes stay. Check shapes first: it catches most bugs before any arithmetic.

Write both shapes (or “doesn’t fit”), then compute $PQ$.

$(3 \times 2)(3 \times 2)$ fails because the inner sizes are 2 and 3. For $PQ$, row 1: $1 \cdot 3 + 2 \cdot 1 = 5$ and $1 \cdot 0 + 2 \cdot 2 = 4$; row 2: $0 \cdot 3 + 1 \cdot 1 = 1$ and $0 \cdot 0 + 1 \cdot 2 = 2$.

$BA$ is $(3 \times 2)(2 \times 3) = 3 \times 3$, not even the same shape as $AB$. Even square matrices usually give different answers in the other order. Matrix multiplication is not commutative.

You may move the brackets but not the matrices. Multiplication is associative. The result is the same either way, but the cost can be very different, as you’ll see in What a matmul costs.

The shapes force the flip: $A^\top$ is $3 \times 2$ and $B^\top$ is $2 \times 3$, so only $B^\top A^\top$ fits to give a $2 \times 2$. It’s the same rule as taking off socks and shoes: the last one on comes off first.

C = A @ B              # matrix product; torch.matmul(A, B) is the same
Q = X @ W_q            # every token's query at once

• $(AB)_{ij} = (\text{row } i \text{ of } A) \cdot (\text{column } j \text{ of } B)$.
• $(m \times k)(k \times n) = (m \times n)$: check shapes first.
• $AB \ne BA$, but $(AB)C = A(BC)$.
• $(AB)^\top = B^\top A^\top$.
• $XW$ applies $W$ to every token (row) at once.