多目的行列（python、numpy）のベクトルの内積

Question

私は2つの行列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 

Pythonのサイクルのために、NとKが大きくなると処理が遅くなることがあります。この乗算を1つにする方法がありますか？numpyの機能はありますか？

Answer 1

合計削減にnp.dotを使用していないようですが、ブロードキャストにつながる拡張のためだけです。したがって、np.newaxis/Noneを使用して1つのディメンションを持つように配列を拡張するだけで、暗黙のブロードキャストが助けになります。

したがって、実装は次のようになります - 新しい軸を追加するための具体的方法を放送上の

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

詳細情報はthis other postで見つけることができます。