Learn AI Create the Future — with Pranay
Main site

PyTorch Tensors

Reductions and the dim argument

sum, mean and argmax all take a dim argument, and the rule nobody tells you is that dim names the dimension that disappears.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

A reduction squeezes many numbers into fewer — a total, an average, a maximum — and dim chooses which direction gets squeezed.

Picture a class register: one row per student, one column per subject. You can total each row to rank students. You can total each column to compare subjects. Same grid, two completely different summaries — the only choice you made was the direction of squeezing.

Why it exists

Almost every question you ask of data is a squeeze. What was the average score? Which product sold most? How wrong was the model, over the whole batch? Each answer is small, made from something big.

With a grid, "add it up" is ambiguous. Down the columns or across the rows? PyTorch needs you to say. The argument that says it is called dim, short for dimension.

How it works

The rule that makes dim finally click: dim names the dimension that disappears.

marks: 2 students x 3 subjects        shape (2, 3)

sum(dim=0)  ->  students disappear  ->  3 subject totals    shape (3,)
sum(dim=1)  ->  subjects disappear  ->  2 student totals    shape (2,)
sum()       ->  everything disappears -> one number          shape ()

Not "the direction it moves along", not "the axis it keeps". The one it removes. Say the sentence "I want one number per subject, so students must disappear, so dim=0" and you will never guess again.

A real example you have seen

Your phone's screen-time report is a stack of reductions. Minutes per app is a squeeze across days. Minutes per day is a squeeze across apps. The weekly total squeezes everything. One grid of app-by-day numbers sits underneath all three views.

Remember this

  • A reduction turns many numbers into fewer: sum, mean, max, argmax.
  • dim names the dimension that disappears from the result.
  • No dim at all squeezes everything down to a single number.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Outputs verified with torch 2.5.1, CPU.

One grid, every direction

reductions.py
import torch

marks = torch.tensor([[80., 90., 70.],
                      [60., 85., 95.]])   # 2 students x 3 subjects

print(marks.sum())            # everything -> one number
print(marks.mean(dim=0))      # dim 0 (students) disappears -> per-subject average
print(marks.mean(dim=1))      # dim 1 (subjects) disappears -> per-student average

print(marks.amax(dim=1))      # each student's best mark
print(marks.argmax(dim=1))    # WHICH subject that best mark was in

total = marks.sum(dim=1, keepdim=True)   # (2, 1) instead of (2,)
print(total.shape)
share = marks / total         # broadcasting lines up because the dim was kept
print(share.round(decimals=2))
Output
tensor(480.)
tensor([70.0000, 87.5000, 82.5000])
tensor([80., 80.])
tensor([90., 95.])
tensor([1, 2])
torch.Size([2, 1])
tensor([[0.3300, 0.3800, 0.2900],
        [0.2500, 0.3500, 0.4000]])

The walkthrough

mean(dim=0) has 3 numbers, mean(dim=1) has 2. Check both against the rule: the named dimension vanished, the other survived. Reading reductions becomes mechanical once you verify shapes this way.

amax versus argmax. amax returns the winning values; argmax returns the winning positions. Student 0's best is 90, sitting at position 1. Classification uses argmax(dim=1) on model outputs constantly: "which class scored highest, per sample".

keepdim=True keeps the squeezed dimension as size 1: shape (2, 1) instead of (2,). That matters the moment you use the result against the original tensor. marks / total needs (2, 3) against (2, 1) — a legal broadcast. Against plain (2,) it would either error or, worse, broadcast in the wrong direction.

Integer tensors refuse mean:

python
torch.tensor([1, 2, 3]).mean()
Output
RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Long

An integer average is usually not an integer, and PyTorch will not round it silently. Cast first: .float().mean().

Common mistakes

Guessing dim by trial and error. Both guesses run without error on a square-ish grid, and one of them is wrong. On a (batch, classes) output, argmax(dim=0) answers "which sample scored highest per class" — nonsense that flows straight into your accuracy metric. Write the sentence: "classes must disappear, so dim=1."

Averaging averages with unequal groups. mean(dim=...) weighs every row equally. If row one has 3 valid entries and row two has 300, the overall mean() and the mean-of-row-means disagree. Decide which you mean before writing either.

Losing the dim you still needed. Dividing by a sum(dim=1) without keepdim is the classic normalisation bug. If a reduction's result will meet the original tensor again, reach for keepdim=True by default.

max versus amax. With a dim, torch.max returns a pair (values, indices), so marks.max(dim=1) needs unpacking. amax returns values alone. Mixing them up produces confusing tuple errors.

Try it yourself

For a (4, 5) tensor of random marks, compute each student's gap from the subject average — original minus per-subject mean — in one expression. The result must stay (4, 5). One of keepdim or broadcasting-from-the-right makes it work; work out which and explain why.

What to learn next

Researcher — Mathematics and papers.

Definition and shape algebra

A reduction along dimension d applies an associative binary operation ⊕ over that index:

out[i_0, …, î_d, …, i_{k−1}] = ⊕_{j=0}^{n_d−1} x[i_0, …, j, …, i_{k−1}]

Where k is the rank, n_d the size of dimension d, and î_d marks the removed index. With keepdim, the output keeps a size-1 dimension at position d, restoring alignment under right-anchored broadcasting — that is its entire purpose.

Cost is O(numel(x)) work regardless of dim; on GPU the practical differences are memory-coalescing (reducing over the fastest-moving stride is cheapest) and the tree-shaped parallel reduction used per output element.

Floating-point non-associativity

Parallel reduction reorders additions, and float addition is not associative: (a + b) + c ≠ a + (b + c) in general. Consequences:

  • The same sum() can differ in the last bits between CPU and GPU, or between tensor shapes, because the reduction tree differs.
  • Naive left-to-right summation of n terms has worst-case relative error O(nε); pairwise/tree summation, which PyTorch's kernels approximate, improves this to O(ε log n) (Higham, 1993, The accuracy of floating point summation).
  • Accumulator dtype matters more than storage dtype: half-precision inputs are typically accumulated in float32 by library kernels. When writing custom reductions, replicate that.

The stable-softmax family

The most important applied reduction is the log-sum-exp appearing in softmax and cross-entropy:

LSE(x) = log Σ_i exp(x_i) = m + log Σ_i exp(x_i − m), with m = max_i x_i

Subtracting the max keeps every exponent ≤ 0, preventing overflow at the cost of nothing — the identity is exact. This is why loss functions that combine softmax and log into a single op (cross_entropy, log_softmax) are both faster and safer than composing them yourself.

argmax subtleties

argmax is piecewise constant, so its derivative is zero almost everywhere — it cannot carry gradients, and PyTorch marks it non-differentiable. Differentiable relaxations exist where a "soft argmax" is genuinely needed: softmax-weighted expectations, and Gumbel-softmax sampling (Jang et al., 2017; Maddison et al., 2017). Tie-breaking is an implementation detail; do not build logic on which duplicate index wins.

What to learn next