Math: Element-wise, Matmul, and Broadcasting

Three categories of math, each with a quirk

Element-wise

Every standard operator (+ - * / **) and most torch.foo functions (exp, log, sqrt, abs, clamp) operate element-wise. The asterisk in a * b is element-wise multiplication, not matrix multiplication — this is the single most common confusion when porting from textbook math.

Matrix multiplication

Matrix multiply is @ (PEP 465 since Python 3.5) or equivalently torch.matmul. Both work for 2D-on-2D and for batched 3D+ tensors where matmul broadcasts over leading dims. torch.bmm is a strictly batched version useful when you want to assert shape contracts.

Reductions

Reductions (sum, mean, max, argmax, std) collapse one or more dims. The dim argument controls which: x.sum() over everything, x.sum(dim=0) down the rows. The default for keepdim is False — the reduced dim disappears. Set keepdim=True when you need to keep broadcasting compatibility downstream.

Broadcasting

Broadcasting lets PyTorch combine tensors of different shapes by virtually expanding size-1 dims. The rule, applied right-to-left:

Each pair of dims must be either equal, or one of them must be 1.
Missing leading dims are treated as 1.

This is identical to NumPy. When in doubt, write down the two shapes one above the other, right-aligned, and check pair-by-pair from the right.

Code

Element-wise vs matrix multiply·python

import torch

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

# ELEMENT-WISE — what the asterisk does
a * b
# tensor([[ 5., 12.],
#         [21., 32.]])

# MATRIX MULTIPLY — what @ does
a @ b
# tensor([[19., 22.],
#         [43., 50.]])

# Both also exist as named functions
torch.mul(a, b)        # element-wise
torch.matmul(a, b)     # matrix multiply

Batched matmul (the attention pattern)·python

import torch

# Q, K, V in attention: (batch, heads, seq, head_dim)
Q = torch.randn(2, 8, 64, 32)
K = torch.randn(2, 8, 64, 32)

# Attention scores: (batch, heads, seq, seq)
# K.transpose(-2, -1) → (2, 8, 32, 64)
scores = (Q @ K.transpose(-2, -1)) / (32 ** 0.5)
print(scores.shape)   # torch.Size([2, 8, 64, 64])

# torch.bmm is the strictly-3D version (no broadcasting on the leading dim)
A = torch.randn(8, 3, 4)
B = torch.randn(8, 4, 5)
torch.bmm(A, B).shape  # torch.Size([8, 3, 5])

Reductions and keepdim·python

import torch

t = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])

t.sum()             # tensor(21.) — over everything
t.sum(dim=0)        # tensor([5., 7., 9.]) — collapse rows → shape (3,)
t.sum(dim=1)        # tensor([6., 15.])     — collapse cols → shape (2,)

# keepdim=True preserves the dim, which keeps broadcasting valid downstream
mean_per_row = t.mean(dim=1, keepdim=True)   # shape (2, 1)
centered = t - mean_per_row                  # broadcasts cleanly
print(centered)

Broadcasting in practice·python

import torch

t = torch.zeros(3, 4)
row = torch.tensor([1, 2, 3, 4])      # (4,)        broadcasts down rows
col = torch.tensor([[10], [20], [30]])  # (3, 1)    broadcasts across cols

(t + row).shape   # torch.Size([3, 4])
(t + col).shape   # torch.Size([3, 4])
(t + row + col).shape  # torch.Size([3, 4])

# Right-aligned shape check:
#         (3, 4)
#            (4,)   ← matches col 4, missing dim treated as 1
#         (3, 1)    ← matches col 4 via 1, row matches 3

Exercise

Implement attention manually for a (batch=2, heads=4, seq=8, head_dim=16) Q, K, V triple. Compute scores, scaled by sqrt(head_dim), softmax over the last dim, then attended values. Compare your output shape against torch.nn.functional.scaled_dot_product_attention (modern PyTorch). They should match exactly.