python - numpy einsum：嵌套点积

Question

我有两个n逐个数k组3和a，b例如，

import numpy as np

a = np.array([
    [
        [1, 2, 3],
        [3, 4, 5]
    ],
    [
        [4, 2, 4],
        [1, 4, 5]
    ]
    ])
b = np.array([
    [
        [3, 1, 5],
        [0, 2, 3]
    ],
    [
        [2, 4, 5],
        [1, 2, 4]
    ]
    ])

它喜欢计算所有“三元组”对的点积，即

np.sum(a*b, axis=2)

一个更好的方法可能是einsum，但我似乎无法弄清楚索引。

这里有什么提示吗？

Answer 1

您正在通过减和减少这两个3D输入数组上的第三个轴，同时保持前两个轴对齐。因此，如果使用np.einsum，我们将使前两个字符串相同，而第三个字符串也相同，但将在输出字符串表示法中跳过，表示我们正在沿该轴为两个输入进行缩减。因此，解决方案是 -

np.einsum('ijk,ijk->ij',a,b)