我如何应用 np.tensordot '沿轴'，例如一个表示样本的轴？

Question

考虑两个多维数组，其中每个数组的第一个轴例如表示不同的样本，这样两个数组在该轴上具有相同的维度，我们希望在这些数组之间沿该维度应用 np.tensordot。

也就是说，我们要实现以下目标：

import numpy as np
a = np.random.normal(size=(10,2,4,5))
b = np.random.normal(size=(10,2))
out = np.zeros((10,4,5))

for i in range(10):
    out[i] = np.tensordot(b[i],a[i],axes=1)

但使用 numpy 操作/索引。

有没有聪明的方法来做到这一点（也有不同的维度，这些只是例子）

Answer 1

我认为einsum 功能可以帮助您解决这个问题

import numpy as np
a = np.random.normal(size=(10,2,4,5))
b = np.random.normal(size=(10,2))
out = np.zeros((10,4,5))

for i in range(10):
    out[i] = np.tensordot(b[i],a[i],axes=1)

out1 = np.einsum('ijkl,ij->ikl', a, b)

(out1 == out).all()

输出：

True