Optimize computation of anomaly from mean along one axis for large arrays

My task is basically a very common one, removing the mean along one axis for a large array (I actually need to compute means and anomalies for subarrays along one dimension). Although this “demeaning/anomaly” calculation is mentioned in a variety of posts, I could not find answers which focus on the efficiency of the operation.

The common way to do this is:

anom = X - X.mean(axis=0)

but as this numpy operation is not parallelized, it just uses one core even if axis1 is very long. Naively I thought this should be optimizable through parallelization? In my main program I have to repeat this kind of operation many times, so saving as much time as possible on this would be great.

I have two dimensional arrays, where the length of axis0, over which I want to compute the mean, is small in comparison to the second axis.

import numpy as np
from numba import njit,prange
X=np.random.randn(20,2000000)

I want to split the axis 0 in some chunks (e.g. length 2,5,10 and 20) and calculate the anomaly and mean over these subblocks (for each element along axis1).

I came up with two functions to do that, the first one just uses numpy and the second implements the task explicitly with for loops using numba with @njit(parallel=True)

def numpy_no_numba(array_in,size):
    s=array_in.shape
    #number of subarrays along axis0
    number=s[0]//size

    mean=np.empty((number,s[1]))
    anom=np.empty_like(array_in)

    #loop over the subarrays
    for i in range(number):
        block=array_in[i*size:(i+1)*size,:]
        block_mean=block.mean(axis=0)
        mean[i]=block_mean
        anom[i*size:(i+1)*size,:]=block-block_mean
    
    return mean,anom

@njit(parallel=True)
def numpy_with_numba(array_in,size):

    s=array_in.shape
    number=s[0]//size
    
    #Initialize array for mean and anomaly (latter has original size)
    mean=np.empty((number,s[1]))
    anom=np.empty_like(array_in)
    
    #looping over blocks if size!=s[0]
    if number>1:
        #loop over all points along axis1
        for i in prange(s[1]):
            for j in prange(number):
                vals=array_in[j*size:(j+1)*size,i]
                m=vals.mean()
                mean[j,i]=m
                anom[j*size:(j+1)*size,i]=vals-m
    else:
        for i in prange(s[1]):
            vals=array_in[:,i]
            m=vals.mean()
            mean[0,i]=m
            anom[:,i]=vals-m
            
    return mean,anom

Timing both methods I get an improvement of a factor of ~3.5:

%%timeit 
m1,a1=numpy_no_numba(X,5)
152 ms ± 73.1 µs per loop 

%%timeit 
m2,a2=numpy_with_numba(X,5)
42 ms ± 468 µs per loop

Are there more efficient ways to solve this? Is this already the best way to use numba, could I actually vectorize the first function better or is gpu acceleration the only way to improve this calculation?

I'm new to parallelization and vectorization in numpy and would really appreciate your input, thanks!