Optimización en tiempo de ejecución en bucle for y malla de rejilla en Python con numpy.

Esta publicación es similar a una pregunta que ya publiqué hace algunas semanas, pero el tiempo de ejecución que tengo actualmente no es muy bueno.
Tengo que implementar esta regla:

<img src="https://render.githubusercontent.com/render/math?math=w_{ij}^{\sigma}=w_{ij}^{\sigma-1}%2b\frac{1}{N}(p_i^{\sigma}p_j^{\sigma}-p_i^{\sigma}h_{ji}^{\sigma}-p_j^{\sigma}h_{ij}^{\sigma}))">

donde

<img src="https://render.githubusercontent.com/render/math?math=h_{ij} = \sum_{k:i\neq k\neq j}w_{ik}^{\sigma-1}p_{k}^{\sigma}">

Terminé con este fragmento de código:

def regla(muestras):

    M, N = muestras.shape
    w = np.zeros((N,N))

    for sigma in range(M):
        H = np.zeros((N,N))

        w_with_diag_zeros = w - np.diag(np.diag(w)) #Exclusión de los casos en los que i == k y forma de (2500, 2500)
        sigma_th_samples_broadcast = muestras[sigma] + np.zeros((N, N)) #Excluir los casos en los que k == j
        muestras_with_diag_zeros = ( sigma_th_samples_broadcast 
                                    - np.diag(np.diag(sigma_th_samples_broadcast)) ) #(2500, 2500)
        #Matriz H
        H = np.matmul(w_with_diag_zeros, np.transpose(muestras_with_diag_zeros))

        #Para la matriz w
        p_i, p_j = np.meshgrid(muestras[sigma], muestras[sigma]) # p_j = p_i también tiene forma (2500, 2500)

        w += (1/N) * (p_i * p_j 
                         - p_j * H.T
                         - p_i * H)  #formas de p_j*H.t = p_i*H = (2500, 2500)

    return w

Las muestras son una matriz 2D de forma (M, N). Es importante destacar que, con M = 50 y N = 2500, crear una matriz con esta función tarda aproximadamente de 40 segundos a 1 minuto. Intenté numpa, cpython, otras definiciones, pero ninguna de estas soluciones proporcionó un mejor resultado. Mi pregunta es, ¿hay alguna manera de optimizar aún más esta función?