多维矩阵中向量的点积（python，numpy）

Question

我有两个矩阵A，B，NxKxD维，我想得到矩阵C，NxKxD维数，其中C[n，k]=A[n，k]xb[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的for-cycle，对于big N和K可能会比较慢。有什么方法可以在一个numpy函数中实现这个乘法吗？在

Answer 1

似乎您并没有使用np.dot来减少和，而只是用于扩展，从而导致广播。因此，您只需使用np.newaxis/None将数组扩展到多个维度，并让隐式广播发挥作用。在

因此，实施将是-

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

关于广播的更多信息，特别是如何添加新轴可以在^{}中找到。在