Acelerar el cálculo de la matriz de distancia con Numpy y Cython

Considere una matriz A numpy de dimensionalidad NxM. El objetivo es calcular la matriz de distancia euclidiana D, donde cada elemento D [i, j] es la distancia de Eucledean entre las filas i y j. ¿Cuál es la forma más rápida de hacerlo? Este no es exactamente el problema que necesito resolver, pero es un buen ejemplo de lo que estoy tratando de hacer (en general, podrían usarse otras métricas de distancia).

Esto es lo más rápido que se me ocurrió hasta ahora:

n = A.shape[0]
D = np.empty((n,n))
for i in range(n):
    D[i] = np.sqrt(np.square(A-A[i]).sum(1))

¿Pero es la forma más rápida? Me preocupa principalmente el ciclo for. ¿Podemos vencer esto con, digamos, Cython?

Para evitar bucles, intenté usar la transmisión y hacer algo como esto:

D = np.sqrt(np.square(A[np.newaxis,:,:]-A[:,np.newaxis,:]).sum(2))

Pero resultó ser una mala idea, porque hay algunos gastos generales en la construcción de una matriz 3D intermedia de dimensionalidad NxNxM, por lo que el rendimiento es peor.

Intenté con Cython. Pero soy un novato en Cython, así que no sé qué tan bueno es mi intento:

def dist(np.ndarray[np.int32_t, ndim=2] A):
    cdef int n = A.shape[0]    
    cdef np.ndarray[np.float64_t, ndim=2] dm = np.empty((n,n), dtype=np.float64)      
    cdef int i = 0    
    for i in range(n):  
        dm[i] = np.sqrt(np.square(A-A[i]).sum(1)).astype(np.float64)              
    return dm 

El código anterior fue un poco más lento que el bucle for de Python. No sé mucho sobre Cython, pero supongo que podría lograr al menos el mismo rendimiento que para loop + numpy. Y me pregunto si es posible lograr una mejora notable en el rendimiento cuando se hace de la manera correcta. ¿O si hay alguna otra forma de acelerar esto (sin involucrar cálculos paralelos)?

Respuestas a la pregunta(1)

Su respuesta a la pregunta