Matrix Multiplication

Matrix multiplication is the operation that dominates deep learning compute. Attention, linear layers, and every projection are built from it. Its shape rule is one line long, and most tensor bugs are violations of it.

The idea

For 2-D tensors, (m, k) @ (k, n) produces (m, n). The inner dimensions must match; the outer dimensions define the output. Entry [i, j] is the dot product of row i on the left with column j on the right.

Three functions do the job with different rules:

• a @ b — operator form, dispatches to torch.matmul.
• torch.matmul(a, b) — supports batching and broadcasting over leading dimensions.
• torch.mm(a, b) — strictly 2-D, no broadcasting.

With batches, leading dimensions broadcast: (B, m, k) @ (k, n) → (B, m, n). With 1-D operands, (k,) @ (k, n) → (n,) and (m, k) @ (k,) → (m,).

torch.nn.Linear(in_features, out_features) stores its weight as (out_features, in_features) and computes x @ W.T + b. That single transpose is why an input of (B, in_features) yields (B, out_features).

Worked example

Project 8 token embeddings of width 64 onto 128 dimensions:

• x → (8, 64)
• W → (128, 64)
• x @ W.T → (8, 128)
• plus bias (128,), broadcast → (8, 128)

Now batched attention scores. Let q and k each be (2, 4, 8, 16) — 2 sequences, 4 heads, 8 tokens, head width 16. Transposing k on its last two axes gives (2, 4, 16, 8), so q @ k.transpose(-2, -1) is (2, 4, 8, 8). The final 8 × 8 block holds one score per token pair, computed independently inside each head.

In code

import torch

x = torch.randn(8, 64)
w = torch.randn(128, 64)
print((x @ w.T).shape)                # torch.Size([8, 128])

q = torch.randn(2, 4, 8, 16)
k = torch.randn(2, 4, 8, 16)
scores = q @ k.transpose(-2, -1)
print(scores.shape)                   # torch.Size([2, 4, 8, 8])

print(torch.mm(torch.randn(3, 5), torch.randn(5, 2)).shape)
# torch.Size([3, 2])
print(torch.matmul(torch.randn(5, 2, 3), torch.randn(3, 7)).shape)
# torch.Size([5, 2, 7])

Check yourself

1. Can you multiply a (4, 3) tensor by another (4, 3) without transposing? Why or why not?
2. What shape does torch.matmul return for (5, 2, 3) @ (3, 7)?
3. In nn.Linear(64, 128), what is weight.shape, and what does x @ weight.T compute?

Key takeaways

• Inner dimensions must match; outer dimensions become the output shape.
• matmul broadcasts leading dimensions, mm is strictly 2-D.
• nn.Linear stores (out, in) and applies x @ W.T + b.