布尔索引为何反转轴顺序？numpy数组x[4,:,i]形状疑问

先看示例代码：

import numpy as np

x = np.zeros((5,4,3,2))
i = np.array([[True, False], [False, False], [False, True]])
print(x[4,:,i].shape) # (2, 4)

疑问点：包含2个True值的索引数组位于:之后，对应长度为4的轴，按预期结果应为(4, 2)才更合理。

解答 #

这是因为numpy的布尔索引遵循提取元素维度前置的规则，而非保留原轴的顺序，具体拆解如下：

• 原数组x形状为(5,4,3,2)，x[4,:,i]先固定第0轴索引为4，此时剩余轴对应原第1、2、3轴，形状为(4,3,2)。
• 布尔数组i形状是(3,2)，对应原第2轴（长度3）和第3轴（长度2）。索引时numpy会遍历i，提取所有True对应的元素，这里共2个匹配项。
• 布尔索引提取的元素会被放在最前面的新维度，剩下未被索引的轴（原第1轴，长度4）紧随其后，所以最终形状是(2,4)。

如果想要得到(4,2)的形状，可以通过调整轴顺序实现，比如：

print(x[4,:,i].transpose(1,0).shape) # (4, 2)

内容的提问来源于stack exchange，提问作者Jeffrey Chen