(2, 3)
tensor reduced over dim=1
becomes (2,)
by default, but (2, 1)
if you set keepdim=True
. The values are identical, but that single 1
in the shape is the difference between a successful normalization and a cryptic RuntimeError
—or worse, a calculation that runs but produces mathematically wrong results.## The mechanics of reduction
In PyTorch, when you use sum
, mean
, max
, or min
, the dimension you specify is the one that vanishes.
import torch
m = torch.tensor([[1., 2., 3.],
[4., 5., 6.]]) # shape (2, 3)
print(m.sum(dim=0).shape) # torch.Size([3]) - rows collapse, columns remain
print(m.sum(dim=1).shape) # torch.Size([2]) - columns collapse, rows remain
If you name a dimension, it's gone. keepdim=True
simply prevents this deletion by leaving a length-1 placeholder.
row_sum = m.sum(dim=1) # shape (2,)
row_sum_k = m.sum(dim=1, keepdim=True) # shape (2, 1)
print(row_sum)
print(row_sum_k)
Why this is a practical tutorial for broadcasting #
Most people use reductions to perform operations back on the original tensor—like row-wise normalization or subtracting a mean. This requires broadcasting. PyTorch aligns shapes from the right; it treats a 1
as "stretch me to fit."
A (2, 1)
tensor broadcasts perfectly against a (2, 3)
tensor because the 1
expands to 3
. A (2,)
tensor does not align.
normed = m / m.sum(dim=1, keepdim=True)
print(normed.sum(dim=1)) # tensor([1., 1.]) - Works perfectly
The "Silent Bug" Trap #
The danger is that forgetting keepdim
fails in two different ways depending on your tensor shape.
Scenario A: The Loud Failure (Non-Square Tensors)
If you have a (2, 3)
tensor and divide by a (2,)
sum, PyTorch tries to align the 3
with the 2
. It crashes immediately with a RuntimeError
. This is actually the best-case scenario because you know exactly where the bug is.
Scenario B: The Silent Failure (Square Tensors)
If your tensor is (3, 3)
, the row sum is (3,)
. PyTorch aligns the rightmost dimensions (3 and 3), and it works—but it does the wrong thing. Instead of dividing each row by its own sum, it divides each column by the row sums. You get a result, but your data is now garbage.
If you're building a real-world AI workflow or a custom LLM agent layer, always default to keepdim=True
for reductions. It removes the ambiguity and prevents these shape-shifting bugs from leaking into your gradients.
Next Why I stopped trusting AI agents to grade their own work →