在Python numpy中进行张量积而不进行求和

Question

我有两个二维的张量，它们有一个共同的长轴（比如说20,000），但是短轴不同，一个是9，另一个是10。我想把它们合并成一个9X10X20000的张量，也就是说在长轴的每一个位置上，其他两个轴是它们的张量积。

具体来说，如果长轴是4，我想做这样的操作：

A = np.arange(8).reshape(2,4)
B = np.arange(12).reshape(3,4)
C = np.zeros(2,3,4)
for i in range(2):
    for j in range(3):
        for k in range(4):
            C[i,j,k] = A[i,k]*B[j,k]

这段代码可以正常工作，但我在想：有没有一种numpy的方法可以做到这一点，而不需要使用循环？

这个问题的背景是为了训练一个神经网络，长轴代表训练样本。当我计算成本函数的梯度时，会得到这样的公式。

谢谢，

Leo

Answer 1

这个例子你也可以用普通的广播来实现：

>>> A[:, np.newaxis, :] * B[np.newaxis, :, :]
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])

上面的代码和下面这个简单的代码是一样的：

>>> A[:, None] * B
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])

Answer 2

我通常会使用 np.einsum 来解决这些问题，这样可以很方便地指定索引应该怎么处理：

>>> A = np.arange(8).reshape(2,4)
>>> B = np.arange(12).reshape(3,4)
>>> np.einsum('ik,jk->ijk', A, B)
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])