PyTorch keepdim: Stop the silent broadcasting bugs

PyTorch's keepdim parameter prevents silent broadcasting bugs when reducing tensors, according to a practical tutorial. Forgetting keepdim can cause either a RuntimeError with non-square tensors or mathematically wrong results with square tensors, as the dimension alignment silently shifts. The tutorial advises always defaulting to keepdim=True for reductions in AI workflows.

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/