Shape Reasoning

Rules for predicting the shape of every operation before you run it.

The four rules

Predict every output shape with four checks: (1) matmul — inner dims must match and cancel, (a, k) @ (k, b) → (a, b); (2) broadcasting — align shapes from the right, each pair of axes must be equal or one of them 1; (3) reduction — a reduce over axis=i removes that axis (or keeps it as 1 with keepdims=True); (4) rank — (n,) is a 1-D vector, not the same as a row (1, n) or column (n, 1).

Matrix multiply & transpose

Operation	Input shapes	Output shape
`A @ B`	`(m, n)`, `(n, p)`	`(m, p)`
`A @ B` (mismatch)	`(m, n)`, `(k, p)` with $n \ne k$	Error — inner dims must match
`A @ x`	`(m, n)`, `(n,)`	`(m,)`
`A.T`	`(m, n)`	`(n, m)`
Batched matmul	`(b, m, n)`, `(b, n, p)`	`(b, m, p)` — leading dims broadcast

Broadcasting & reductions

Operation	Input shapes	Output shape
Add row vector	`(m, n)` + `(n,)`	`(m, n)`
Add column vector	`(m, n)` + `(m, 1)`	`(m, n)`
Outer via broadcast	`(m, 1)` * `(1, n)`	`(m, n)`
Incompatible	`(m, n)` + `(m,)` with $n \ne m$	Error — right-aligned axes clash
`x.sum(axis=0)`	`(m, n)`	`(n,)`
`x.sum(axis=1, keepdims=True)`	`(m, n)`	`(m, 1)`
`x.sum()`	`(m, n)`	`()` (scalar)