dot произведение векторов в многомерных матрицах (python, numpy)

Question

У меня есть две матрицы A, B, NxKxD, и я хочу получить размерность матрицы C, NxKxDxD, где C [n, k] = A [n, k] x B [n, k] .T (здесь «x» означает произведение матриц размерностей Dx1 и 1xD, поэтому результат должен быть DxD мерным), поэтому теперь мой код выглядит следующим образом (здесь A = B = X):

def square(X): 
    out = np.zeros((N, K, D, D)) 
    for n in range(N): 
     for k in range(K): 
      out[n, k] = np.dot(X[n, k, :, np.newaxis], X[n, k, np.newaxis, :]) 
    return out 

Это может быть медленным для больших N и K из-за цикла python. Есть ли способ сделать это умножение в одном numpy функция?

Answer 1

Кажется, вы не используете np.dot для суммирования, но только для расширения, что приводит к трансляции. Таким образом, вы можете просто расширить массив, чтобы иметь еще одно измерение с использованием np.newaxis/None и позволить скрытому трансляции помочь.

Таким образом, реализация будет -

X[...,None]*X[...,None,:] 

Более подробная информация о вещании в частности, как добавить новые оси может быть найдена в this other post.