Per-channel statistics

A batch of images is (batch, channels, height, width), and you want the mean of each channel across the whole batch: one number per channel.

Three axes have to disappear and one has to stay.

x.mean(dim=(0, 2, 3))       # (channels,)

dim takes a tuple. Every axis in it is reduced, and what is left is what was not mentioned.

Why not reduce one at a time

You can:

x.mean(dim=3).mean(dim=2).mean(dim=0)

and it gives the same answer here, but only because every axis has the same weight in a mean over a rectangular tensor. Two things make it worse code:

• The indices shift. Each reduction removes an axis, so the numbers you pass afterwards refer to different axes than they did before. Reducing in the other order needs different indices for the same result, and getting that wrong is silent.
• It is three passes. Three kernels over progressively smaller data, where one would do.

For a sum the answers agree; for mean they agree here and would not if the reductions were weighted; for max they agree. The habit of naming all the axes at once is right regardless, because it says what you meant rather than a sequence that happens to arrive there.

Negative indices help

dim=(0, 2, 3) breaks if a rank changes. dim=(0, -2, -1) says “the batch and the two spatial axes” and survives.

keepdim applies to all of them

x.mean(dim=(0, 2, 3), keepdim=True)     # (1, channels, 1, 1)

which is exactly the shape you need to normalise x by it, and is what BatchNorm holds.

Your task

def channel_means(x: torch.Tensor) -> torch.Tensor

x is (batch, channels, height, width). Return a 1-D tensor of length channels: the mean over everything else. Use at most two dispatched operations.