Сумма над квадратом массива

Question

как часть пакетного евклидово расстояние вычисления, я вычисления

(X * X).sum(axis=1) 

, где X является достаточно большой 2-D массив. Это прекрасно работает, но он создает временный массив того же размера, что и X. Есть ли способ избавиться от этого временного, но сохранить эффективность векторной операции?

Очевидным кандидатом,

np.array([np.dot(row, row) for row in X]) 

работает, но использует список Python в качестве временного, что делает его довольно медленно.

Без axis, память-эффективная форма будет

(X * X).sum() => np.dot(X.ravel(), X.ravel()) 

, и я знаю, что, когда axis=1, это эквивалентно

np.diag(np.dot(X, X.T)) 

который заставил меня смотреть в обобщениях dot, таких как np.inner, np.tensordot и np.einsum, но я не могу понять, как они решат мою проблему.

Answer 1

einsum эквивалент:

np.einsum('ij,ij->i', X, X) 

Хотя я не уверен, как это работает внутри, так что может или не может не решить вашу проблему.