python - numpy einsum 获取轴置换

Question

我在“np.einsum”的文档中了解到，排列字符串会给出向量中轴的排列。以下实验证实了这一点:

>>> M = np.arange(24).reshape(2,3,4)
>>> M.shape
(2, 3, 4)
>>> np.einsum('ijk', M).shape
(2, 3, 4)
>>> np.einsum('ikj', M).shape
(2, 4, 3)
>>> np.einsum('jik', M).shape
(3, 2, 4)

但是这个我无法理解:

>>> np.einsum('kij', M).shape
(3, 4, 2)

我会期望 (4, 2, 3) 而不是...我的理解有什么问题吗？

Answer 1

当未指定输出签名时(即下标字符串中没有 '->')，einsum 将通过获取给定的字母来创建它，并且按字母顺序排列它们。

这意味着

np.einsum('kij', M)

实际上等同于

np.einsum('kij->ijk', M)

因此，编写 'kij' 标记了输入 矩阵的轴，而不是输出矩阵，这会导致您观察到的轴排列。

这一点在文档中没有明确说明，但可以在 C source code for einsum 中看到注释:

/*
 * If there is no output signature, create one using each label
 * that appeared once, in alphabetical order
 */

为确保 M 的轴按预期顺序排列，可能需要为输入和输出矩阵赋予 einsum 标签:

>>> np.einsum('ijk->kij', M).shape
(4, 2, 3)