How to get numpy to broadcast an operation after a reduction operation

Question

I am trying to normalize some data for the last dimensions.

#sample data
x = numpy.random.random((3, 1, 4, 16, 16))
x[1] = x[1]*2
x[2] = x[2]*4

I can get the mean,

m = x.mean((-3, -2, -1))

Now, x.shape is (3, 1, 4, 16, 16) and m.shape is (3, 1), I want to subtract the mean from each sample. So far I have.

for i in range(x.shape[0]):
    for j in range(x.shape[1]):
        x[i,j] = x[i,j] - m[i,j]

That works, but it has two drawbacks. I'm using explicit loops, and it it requires the shape to have 5 dimensions.

Answer 1

Simply keep dimensions with keepdims arg and then subtract -

m = x.mean((-3, -2, -1),keepdims=True)
x -= m

This would work regardless of the axes that are used for the reduction and should be a clean solution.