PYTHON / NUMPY
Aggregations and axis semantics
Predict and control the shape of any NumPy reduction by treating axis as the dimension that gets collapsed, and use keepdims to broadcast results back.
What you will learn
- Predict a reduction's shape by deleting the collapsed axis from the input shape
- Use axis=1 for per-row results and axis=0 for per-column results on 2D arrays
- Add keepdims=True so a reduction still broadcasts against the original array
- Collapse several dimensions at once with a tuple, such as axis=(0, 2)
Understanding Aggregations and axis semantics
An aggregation like sum, mean, max or std removes one index from an array. The axis argument names which index position disappears, not the direction you visually scan. For a 2D array a, a.sum(axis=1) produces out[i] = a[i, 0] + a[i, 1] + ... because the second index is the one being summed over, so the (3, 4) input becomes a (4,)-free result of shape (3,). The rule generalises to any rank: take the input shape, delete the entry at position axis, and that is the output shape.
That deleted entry is exactly what breaks the next line of your code. Broadcasting aligns shapes from the right, so a (3,) row-mean no longer lines up with the (3, 4) array it came from. Passing keepdims=True leaves a length-1 axis in place, giving (3, 1), which broadcasts back across the rows and lets you write a - a.mean(axis=1, keepdims=True). Passing a tuple collapses several axes in one call, and the default axis=None collapses every axis down to a scalar.
The reduction also decides the dtype and the edge cases. mean on an integer array always returns float64, sum of a boolean array counts True values, and sum over a zero-length axis returns 0 because addition has an identity while max raises ValueError because there is no smallest value to fall back on. A single nan poisons the whole reduction it participates in, which is why nansum and nanmean exist, and argmax returns a position along the collapsed axis rather than the value found there.
import numpy as np
a = np.arange(12).reshape(3, 4)
print(a)
print("shape:", a.shape)
print("sum of everything:", a.sum())
print("axis=0 collapses rows:", a.sum(axis=0), a.sum(axis=0).shape)
print("axis=1 collapses columns:", a.sum(axis=1), a.sum(axis=1).shape)
print("keepdims keeps the slot:", a.sum(axis=1, keepdims=True).shape)
print(a - a.mean(axis=1, keepdims=True))The axis argument names the dimension that the aggregation consumes, so the result's shape is the input's shape with that one entry removed.
Worked examples
Axis tuples and negative axes in 3D
Shows how the shape rule extends past two dimensions and how a tuple collapses more than one axis.
import numpy as np
v = np.arange(24).reshape(2, 3, 4)
print(v.sum(axis=0).shape)
print(v.sum(axis=(0, 2)))
print(v.sum(axis=-1).shape)
print(v.max(axis=1))Example explained
Line 1v.sum(axis=0) deletes the first entry of (2, 3, 4), leaving (3, 4).
Line 2axis=(0, 2) deletes both of those entries at once, so only the length-3 middle axis survives.
Line 3axis=-1 is the last axis whatever the rank, so (2, 3, 4) becomes (2, 3).
Line 4v.max(axis=1) compares the three middle slices elementwise, keeping the largest, which is the j=2 slice.
argmax returns a position, not a value
Contrasts per-axis argmax with the flat index you get when no axis is given.
import numpy as np
temps = np.array([[12, 19, 15, 11],
[18, 14, 21, 17],
[ 9, 25, 13, 22]])
print(temps.argmax(axis=1))
print(temps.max(axis=1))
print(temps.argmax())
print(np.unravel_index(temps.argmax(), temps.shape))Example explained
Line 1argmax(axis=1) collapses the columns and reports the column index of each row's largest value.
Line 2max(axis=1) collapses the same axis but reports the values 19, 21 and 25 instead of 1, 2 and 1.
Line 3With no axis the array is treated as flat, so 25 is reported as index 9, not as (2, 1).
Line 4np.unravel_index converts that flat index back into row and column coordinates.
Why keepdims exists
Demonstrates the broadcast failure you get when a collapsed axis is not kept, and the fix.
import numpy as np
counts = np.array([[2, 3, 5],
[1, 1, 8]], dtype=float)
totals = counts.sum(axis=1)
print(totals.shape)
try:
print(counts / totals)
except ValueError:
print("cannot divide (2, 3) by (2,)")
totals2 = counts.sum(axis=1, keepdims=True)
print(totals2.shape)
print(counts / totals2)Example explained
Line 1Summing axis=1 of a (2, 3) array gives shape (2,), one total per row.
Line 2Broadcasting pads on the left, so (2,) is compared against the length-3 last axis and 2 != 3 fails.
Line 3keepdims=True gives (2, 1), whose length-1 axis stretches across the three columns.
Line 4Each row now divides by its own total, so every row of the result sums to 1.
Important notes
mean always returns float64 for integer input, while sum keeps an integer dtype and upcasts small types such as int8 to the platform integer, so check the dtype when totals matter.
axis=-1 means the last axis regardless of rank, which keeps per-row code working when you later add leading batch dimensions.
Common mistakes
Reading axis=0 as 'work on the rows' and calling a.sum(axis=0) for row totals: you get one number per column instead, and on a non-square array the result even has the wrong length.
Omitting keepdims before arithmetic: on a square array the (n,) result still broadcasts, so it silently subtracts row statistics along the columns and produces wrong numbers with no error.
Assuming every reduction tolerates emptiness or nan: sum over a zero-length axis returns 0, but max raises ValueError, and one nan makes mean and max return nan unless you switch to nanmean or nanmax.
Try it yourself
Change, predict, then run
Build m = np.arange(12).reshape(4, 3) and print m.sum(axis=0), m.max(axis=1) and their shapes. Then produce a version of m whose columns each sum to 1 using keepdims, and verify with .sum(axis=0).
Open the Python workspaceCheck your understanding
x has shape (3, 3) and you write x - x.mean(axis=1) intending to center each row. What happens?
- It raises a ValueError because (3,) cannot broadcast against (3, 3)
- It runs without error but subtracts the row means across the columns, giving wrong numbers
- It centers each row correctly, since axis=1 already selected the rows
- It returns a one-dimensional array of three centered means
Show answer
x.mean(axis=1) has shape (3,), and broadcasting aligns it with the last axis of x, so each row mean is subtracted from a column instead of a row. Because the array is square the lengths happen to match, so no ValueError is raised, which is what makes this bug hard to spot; on a (3, 4) array the same code would fail loudly. keepdims=True gives (3, 1) and centers the rows as intended.