python - numpy 中的特定张量积

Question

我想在矩阵上计算以下运算:

import numpy as np

x = np.arange(9).reshape((3,3))
result = np.zeros((3,3,3))
for i in range(3):
    for j in range(3):
        for k in range(3):
            result[i,j,k] = x[j,i] * x[j,k]

这给出

array([[[  0.,   0.,   0.],
        [  9.,  12.,  15.],
        [ 36.,  42.,  48.]],

       [[  0.,   1.,   2.],
        [ 12.,  16.,  20.],
        [ 42.,  49.,  56.]],

       [[  0.,   2.,   4.],
        [ 15.,  20.,  25.],
        [ 48.,  56.,  64.]]])

正如预期的那样。

问题

如何使用 numpy 使用张量积(无循环)执行此计算？

编辑

如果 X 的元素是向量，则运算为:

result[i,j,k] = np.dot(x[j,i] , x[j,k])

适合此计算的 numpy 运算符是什么？

Answer 1

使用迭代器作为字符串表达式的直接方式 np.einsum将是-

np.einsum('ji,jk->ijk',x,x)

另一个带有 broadcasting并交换轴 -

(x[:,None,:]*x[:,:,None]).swapaxes(0,1)