Subtract each row's mean with correct broadcasting

Implement centre_rows(x) for a two-dimensional floating-point tensor. Subtract each row's mean from every value in that row. Return a tensor with the same shape as x.

For example:

input:   [[1., 3.], [10., 20.]]
means:   [2., 15.]
output:  [[-1., 1.], [-5., 5.]]

The broadcasting trap

The starter uses:

x - x.mean(dim=1)

For an input shaped (rows, cols), the mean has shape (rows,). Broadcasting aligns dimensions from the right, so that one-dimensional result aligns with the column axis of x.

On a square tensor, the sizes match and the subtraction runs. It subtracts mean number j from column j, instead of subtracting each row's own mean. For torch.arange(16).reshape(4, 4).float(), the first row shows the difference:

required row 0: [-1.5, -0.5,  0.5,   1.5]
starter row 0:  [-1.5, -4.5, -7.5, -10.5]

Compatible shapes do not guarantee correct values. A square case with equal row means can even hide this bug. On shape (3, 4), the starter raises a shape error; on (3, 1), it broadcasts to the wrong output shape (3, 3).

Keep the reduced axis

Using keepdim=True gives the means shape (rows, 1):

x:         (rows, cols)
row means: (rows,    1)

The trailing axis then broadcasts across each row's columns. This preserves the intended pairing for square and rectangular inputs, including a single row or a single column.

Match shapes from the right

Broadcasting lets an operation apply across compatible shapes without explicitly constructing repeated inputs. Compare dimensions from the right: each pair must be equal, or one must be 1. Missing leading dimensions behave like dimensions of 1. At each position, equal sizes stay the same. If one size is 1, the result uses the other size. This also covers empty dimensions: (0, 3) and (1, 3) broadcast to (0, 3), not (1, 3).

For example, (4, 3) plus (3,) applies the same three-feature bias to all four rows. Adding (4, 1) instead applies a different scalar to each row. A size-one axis lets the same value be reused along that axis.

This explains a subtle centring bug: reducing (B, D) over features gives (B,). Subtracting it directly aligns B with the feature axis. If B == D, the operation runs but uses the wrong means. Retaining the reduction axis as (B, 1) expresses the intended dependence.

Check how much memory the result needs

The input need not be physically copied, but the output can still be huge: (N, 1, D) minus (1, M, D) has shape (N, M, D). An in-place operation must fit its result in the existing destination shape. Explicit expansion may also create a view whose positions refer to the same storage, making independent writes unsafe.

Before implementing an operation, write the aligned shapes and the result shape. Then check which inputs are reused and how much space the output needs.