The bug that only shows on a square

Two problems in, you could be forgiven for thinking this course only grades cost. It does not. Most problems here grade the ordinary thing: is the answer right.

What is unusual is which wrong answers get chosen. Not the ones that raise an exception, because those find themselves. The ones that return a perfectly-shaped tensor full of wrong numbers.

Centre each row

Subtract each row’s mean from that row. The obvious line:

x - x.mean(dim=1)

x is (rows, cols). x.mean(dim=1) collapses the column axis and gives (rows,) — dim 1 is gone, which is what reducing over it means.

Now broadcasting aligns the two shapes from the right:

    x          (rows, cols)
    mean             (rows,)
    aligned to  (rows, cols)
                      ^^^^
                      matched against `rows`, not `cols`

Each row’s mean lines up with a column. The subtraction runs, the result has the right shape, and every number is wrong.

Why it survives testing

On a (4, 4) input, rows == cols, so the shapes are compatible and nothing complains:

correct row 0:  [-1.5, -0.5,  0.5,  1.5]
what you get:   [-1.5, -4.5, -7.5, -10.5]

On a (3, 4) input it finally raises:

RuntimeError: The size of tensor a (4) must match the size of
tensor b (3) at non-singleton dimension 1

So this bug is invisible for exactly as long as your test data is square, and square test data is what everybody writes first. This course will keep putting a square case and a rectangular case in front of you for that reason.

The fix

keepdim=True leaves the reduced axis in place with length 1:

x.mean(dim=1)                 ->  (rows,)
x.mean(dim=1, keepdim=True)   ->  (rows, 1)

A length-1 axis broadcasts against anything, and now it is the column axis that stretches, which is what “each row’s mean” meant all along.

Your task

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

Return x with each row’s mean subtracted from that row.

The starter is the obvious line. Read the shapes before you change anything.