python - Numpy einsum 沿轴计算外积

Question

我有两个包含兼容矩阵的 numpy 数组，想计算使用 numpy.einsum 的元素明智的外积.数组的形状是:

A1 = (i,j,k)
A2 = (i,k,j) 

因此数组分别包含 i 个形状为 (k,j) 和 (j,k) 的矩阵。

所以给定 A1 将包含矩阵 A,B,C 并且 A2 将包含矩阵 D,E,F，结果将是:

A3 = (A(x)D,B(x)E,C(x)F)

(x) 是外积运算符。

根据 this answer 我的理解是这样的以下形状的数组 A3:

A3 = (i,j*k,j*k)

到目前为止我已经尝试过:

np.einsum("ijk, ilm -> ijklm", A1, A2)

但是生成的形状不正确。

作为完整性检查，我正在对此进行测试:

A = np.asarray(([1,2],[3,4]))
B = np.asarray(([5,6],[7,8]))

AB_outer = np.outer(A,B)

A_vec = np.asarray((A,A))
B_vec = np.asarray((B,B))

# this line is not correct
AB_vec = np.einsum("ijk, ilm -> ijklm", A_vec,B_vec)

np.testing.assert_array_equal(AB_outer, AB_vec[0])

这目前会引发断言错误，因为我的 einsum 表示法不正确。我也乐于接受任何可以解决这个问题并且比 nymphs einsum 更快或同样快的建议。

Answer 1

我们可以扩展 dims 并让 broadcasting为我们完成工作 -

(A1[:,:,None,:,None]*A2[:,None,:,None,:]).swapaxes(2,3)

sample 运行-

In [46]: A1 = np.random.rand(3,4,4)
    ...: A2 = np.random.rand(3,4,4)

In [47]: out = (A1[:,:,None,:,None]*A2[:,None,:,None,:]).swapaxes(2,3)

In [48]: np.allclose(np.multiply.outer(A1[0],A2[0]), out[0])
Out[48]: True

In [49]: np.allclose(np.multiply.outer(A1[1],A2[1]), out[1])
Out[49]: True

In [50]: np.allclose(np.multiply.outer(A1[2],A2[2]), out[2])
Out[50]: True

相当于np.einsum会是——

np.einsum('ijk,ilm->ijklm',A1,A2)