What Does `keepdim` Do in PyTorch? (And the Silent Bug When You Forget It)

You wrote x.sum(dim=1) to get a row total, divided the original tensor by it, and either got RuntimeError: The size of tensor a (3) must match the size of tensor b (2) — or, worse, got no error at all and numbers that were quietly wrong. Someone told you to add keepdim=True and it fixed itself. But nobody said why.

The short answer

Every reduction — sum, mean, max, min — collapses a dimension, and by default that dimension disappears from the shape. keepdim=True tells PyTorch to keep it as a size-1 axis instead. A (2, 3) tensor reduced over dim=1 becomes (2,) normally, but (2, 1) with keepdim=True. Same numbers either way — the only difference is a single 1 in the shape. That 1 is what lets the result broadcast back against the original tensor.

import torch

m = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.]])

print(m.sum(dim=1).shape)                 # torch.Size([2])
print(m.sum(dim=1, keepdim=True).shape)   # torch.Size([2, 1])

That's the whole feature. The interesting part is why the missing 1 breaks things — sometimes loudly, sometimes silently.

First: what does dim do in a reduction?

Before keepdim makes sense, you need the mental model for dim. The dimension you name is the one that disappears.

m = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.]])   # shape (2, 3)

print(m.sum())          # tensor(21.)          — no dim: collapse everything to a scalar
print(m.sum(dim=0))     # tensor([5., 7., 9.]) — collapse the rows, shape (3,)
print(m.sum(dim=1))     # tensor([ 6., 15.])   — collapse the columns, shape (2,)

dim=0 adds down each column (the size-2 axis vanishes, leaving (3,)). dim=1 adds across each row (the size-3 axis vanishes, leaving (2,)). If you can predict the output shape before running, you understand reductions. The dim you name is gone.

So what does keepdim actually do?

By default the collapsed axis is deleted. keepdim=True leaves it behind as 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.]])

Identical values. The keepdim version is just wrapped so each row sum sits in its own row. That shape — (2, 1) instead of (2,) — is the entire point, because of what happens next.

Why size-1 matters: broadcasting back against the original

The most common reason to reduce a tensor is to then use that result on the original — normalize each row to sum to 1, subtract each column's mean, divide by a per-row max. Every one of those is original (op) reduction, and for it to line up, the reduction has to broadcast.

PyTorch broadcasting aligns shapes from the right and treats a 1 as "stretch me to fit." A (2, 1) sum broadcasts cleanly against a (2, 3) tensor: the 1 stretches across the 3 columns. A bare (2,) sum does not line up at all.

# Correct: keepdim=True → (2, 1) broadcasts against (2, 3)
normed = m / m.sum(dim=1, keepdim=True)
print(normed)
# tensor([[0.1667, 0.3333, 0.5000],
#         [0.2667, 0.3333, 0.4000]])
print(normed.sum(dim=1))   # tensor([1., 1.]) — every row sums to 1 ✅

Drop the keepdim and this exact line either crashes or lies to you. Which one you get depends on the shape — and that's the part almost every tutorial skips.

The trap: the same bug fails two completely different ways

Case 1 — a non-square tensor: it errors (the lucky case)

Our m is (2, 3). Without keepdim, the row sum is (2,), and PyTorch tries to align (2, 3) against (2,) from the right: 3 vs 2. Mismatch.

m = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.]])          # (2, 3)

normed = m / m.sum(dim=1)                  # (2, 3) / (2,)
# RuntimeError: The size of tensor a (3) must match the size of
# tensor b (2) at non-singleton dimension 1

This is the good outcome. PyTorch stops you immediately, and the error text even names the two sizes that don't match. (If that message looks familiar, it's the same shape-mismatch error covered in PyTorch Broadcasting Explained.)

Case 2 — a square tensor: it runs, and it's silently wrong (the dangerous case)

Now make the tensor square. The row-sum vector is length 3, the tensor has 3 columns, so the shapes happen to broadcast — just not the way you meant.

m = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.],
                  [7., 8., 9.]])          # (3, 3)

wrong   = m / m.sum(dim=1)                 # (3, 3) / (3,)  — no error!
correct = m / m.sum(dim=1, keepdim=True)   # (3, 3) / (3, 1)

print(wrong.sum(dim=1))    # tensor([0.4250, 1.2500, 2.0750])  ← NOT 1s ❌
print(correct.sum(dim=1))  # tensor([1., 1., 1.])              ← correct ✅

No exception. No warning. The row sums aren't 1 because the (3,) vector aligned to the columns instead of the rows — PyTorch divided the wrong way. Your data looks normalized, your code ran clean, and the mistake surfaces as a mysteriously bad model three hours later.

This is why you add keepdim=True by reflex, not by debugging. On a square batch it's the difference between correct and quietly-corrupted, with nothing on screen to tell you which one you got.

A rule you can actually remember

Whenever you reduce a tensor and then use the result in arithmetic with the original, pass keepdim=True.

That's it. If the reduction is the final answer (you just want the row totals and you're done), you don't need it. The moment the reduction feeds back into the original tensor, you do.

Does keepdim work on max, min, and argmax too?

Yes — max(dim) and min(dim) return two tensors (the values and their indices), and keepdim keeps the axis on both:

vals, idxs = m.max(dim=1, keepdim=True)
print(vals.shape, idxs.shape)   # torch.Size([3, 1]) torch.Size([3, 1])

This is handy for a "divide each row by its own max" style normalization, where you need that (3, 1) shape to broadcast. argmax(dim, keepdim=True) works the same way.

Common mistakes and gotchas

• Forgetting keepdim on a square tensor. The headline trap: it broadcasts to the wrong axis with no error. Rows-or-columns bugs on square data are almost always a missing keepdim.
• Thinking keepdim changes the numbers. It never does. The values are identical; only the shape gains a 1. If your totals changed, something else is wrong.
• Reaching for unsqueeze after the fact. x.sum(dim=1).unsqueeze(1) gives the same (N, 1) shape, but x.sum(dim=1, keepdim=True) says it in one step and can't get the axis position wrong.
• Assuming "it ran" means it was right. As Case 2 shows, running successfully is not the same as broadcasting correctly. Print .shape on the reduction before you divide.
• Using keepdim when you don't need it. If the reduced value is your final output (not fed back into the original tensor), the trailing 1 is just clutter — drop it.

Recap

keepdim=True keeps a reduced dimension as a size-1 axis instead of deleting it, so the result still broadcasts against the original tensor. Forget it and you get one of two failures: a shape-mismatch RuntimeError on non-square tensors (annoying but safe), or a silent wrong answer on square tensors (dangerous). The habit that saves you: any time you reduce and then combine with the original — normalizing rows, centering columns, scaling by a per-row max — pass keepdim=True.

---

This is one idea from my book PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped, you can grab the free PyTorch tensor cheat-sheet here, run every example from the book in the companion notebooks on GitHub, get the digital edition on Leanpub, or find the full paperback on Amazon here.