Per-channel statistics

For an image batch shaped (batch, channels, height, width), a per-channel mean averages over batch and spatial positions. The channel axis remains:

import torch

x = torch.arange(16, dtype=torch.float32).reshape(2, 2, 2, 2)
y = x.mean(dim=(0, 2, 3))
print(list(y.shape))  # [2]
print(y.tolist())  # [5.5, 9.5]

The tuple names the original axes to reduce. Axis 1, channels, is absent from the tuple and therefore remains in the result.

Why name the axes together?

The starter chains mean(dim=3), mean(dim=2) and mean(dim=0). For this rectangular input, it computes the same mathematical mean: each group at a given stage contains the same number of entries. Floating-point rounding can still differ between implementations.

A tuple makes the intended reduction easier to inspect and avoids tracking how axis numbers change after each step. A wrong index may reduce the wrong axis or raise an error, depending on the remaining rank.

The harness counts dispatched operations. The chained version dispatches three means; the tuple version meets the limit with one mean operation. This count is not a measurement of GPU kernels, memory passes or elapsed time. Benchmark those separately when performance matters.

Keep dimensions for broadcasting

kept = x.mean(dim=(0, 2, 3), keepdim=True)
print(list(kept.shape))  # [1, 2, 1, 1]
centered = x - kept
print(list(centered.shape))  # [2, 2, 2, 2]

Each reduced axis stays present with size one, so the means broadcast back over x. Subtracting them centers each channel; a complete normalization may require additional statistics and scaling. BatchNorm's stored running means are 1-D channel vectors, not this four-dimensional broadcast shape.

Negative indices can name axes relative to the end. For this 4-D input, (0, -2, -1) identifies the same axes as (0, 2, 3). It is not a general solution for arbitrary ranks: extra spatial axes or a different batch layout need a corresponding reduction rule.

Give each tensor axis a meaning

A tensor's shape tells you how many entries each axis has. It does not tell you what those entries represent. For example, (batch, features) and (features, batch) use different axis orders. Name the axes before choosing an operation.

A reduction combines entries along an axis, for example by summing or averaging. Averaging the features of a (B, D) tensor gives one number per example, with shape (B,). Keeping the reduced axis gives shape (B, 1) instead.

For [[2, 4, 6], [1, 3, 5]], the row means with the axis kept are [[4], [3]]. Subtracting them from the input gives [[-2, 0, 2], [-2, 0, 2]]: each row has been centered using its own mean.

Reshape and transpose do different jobs

A reshape regroups elements using the operation's traversal order. A transpose changes the axis order. For the usual row-major reshape, [[1, 2, 3], [4, 5, 6]] reshaped to (3, 2) becomes [[1, 2], [3, 4], [5, 6]]. Transposing the original instead gives [[1, 4], [2, 5], [3, 6]].

Both results have shape (3, 2), but their entries mean different things. Check values as well as shapes. Unequal axis sizes make many axis mistakes easier to spot than square examples do.

Concatenate along an axis or stack on a new one

Concatenation extends an existing axis. Stacking creates a new axis. For two arrays of shape (2, 3), concatenating on axis zero gives (4, 3). Stacking on a new leading axis gives (2, 2, 3).