Espera-se que o código de distância euclidiana quadrada eficiente da numba seja mais lento que o equivalente eficiente da numpy?
Modifico o código mais eficiente de (Por que esse código numba é 6x mais lento que o código numpy?) para que ele possa lidar com x1 sendo (n, m)
@nb.njit(fastmath=True,parallel=True)
def euclidean_distance_square_numba_v5(x1, x2):
res = np.empty((x1.shape[0], x2.shape[0]), dtype=x2.dtype)
for a_idx in nb.prange(x1.shape[0]):
for o_idx in range(x2.shape[0]):
val = 0.
for i_idx in range(x2.shape[1]):
tmp = x1[a_idx, i_idx] - x2[o_idx, i_idx]
val += tmp * tmp
res[a_idx, o_idx] = val
return res
No entanto, ainda não é mais eficiente que a versão do numpy mais eficiente:
def euclidean_distance_square_einsum(x1, x2):
return np.einsum('ij,ij->i', x1, x1)[:, np.newaxis] + np.einsum('ij,ij->i', x2, x2) - 2*np.dot(x1, x2.T)
Com entrada como
a = np.zeros((1000000,512), dtype=np.float32)
b = np.zeros((100, 512), dtype=np.float32)
O tempo que obtive é 2,4723422527313232 para o código numba e 0,8260958194732666 para o código numpy.