python - 如何理解这个 np.einsum('kij',A) 的结果？

Question

例如，

A = np.arange(24).reshape((2, 3, 4))
print np.einsum('ijk', A)

这仍然A没有问题。

但如果我这样做print np.einsum('kij', A)的形状是(3, 4, 2)。不应该(4, 2, 3)吗？

print np.einsum('cab', A)形状的结果也(4, 2, 3)没有问题。为什么print np.einsum('kij', A)不一样？

Answer 1

如果您仅指定一组下标，则这些下标将被解释为输入数组中相对于输出的维度顺序，反之亦然。

例如：

import numpy as np

A = np.arange(24).reshape((2, 3, 4))
B = np.einsum('kij', A)

i, j, k = np.indices(B.shape)

print(np.all(B[i, j, k] == A[k, i, j]))
# True

正如@hpaulj在评论中指出的那样，您可以通过指定两组下标来使输入和输出维度之间的对应关系更加明确：

# this is equivalent to np.einsum('kij', A)
print(np.einsum('kij->ijk', A).shape)
# (3, 4, 2)

# this is the behavior you are expecting
print(np.einsum('ijk->kij', A).shape)
# (4, 2, 3)