# PyTorch keepdim: Stop the silent broadcasting bugs

> Source: <https://promptcube3.com/en/threads/2745/>
> Published: 2026-07-24 13:46:18+00:00

# PyTorch keepdim: Stop the silent broadcasting bugs

`(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.

``` python
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.

[Next Why I stopped trusting AI agents to grade their own work →](/en/threads/2731/)