PyTorch keepdim: Stop the silent broadcasting bugs

A (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) 
# tensor([ 6., 15.])
print(row_sum_k) 
# tensor([[ 6.],
# [15.]])

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.

# Correct: keepdim=True allows (2, 1) to broadcast against (2, 3)
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.