如何沿批量维度广播 numpy 索引？

【问题标题】：How to broadcast numpy indexing along batch dimensions?如何沿批量维度广播 numpy 索引？
【发布时间】：2020-01-25 18:12:28
【问题描述】：

例如，np.array([[1,2],[3,4]])[np.triu_indices(2)] 的形状为 (3,)，是上三角条目的扁平列表。但是，如果我有一批 2x2 矩阵：

foo = np.repeat(np.array([[[1,2],[3,4]]]), 30, axis=0)

我想获得每个矩阵的上三角指数，天真的尝试是：

foo[:,np.triu_indices(2)]

然而，这个对象实际上是 (30,2,3,2) 的形状（与我们可能期望的 (30,3) 不同，如果我们已经批量提取了上三角条目。

我们如何才能沿批次维度广播元组索引？

【问题讨论】：

你所说的“每个矩阵”是什么意思？您只有一个甚至不是方形的矩阵！如果你想要每个 2x2 子数组（轴 1 和 2）的上三角形，你可以这样做 x, y = np.triu_indices(2); foo[:,x,y]。

【解决方案1】：

获取元组并使用它们来索引最后两个维度 -

r,c = np.triu_indices(2)
out = foo[:,r,c]

或者，使用 Ellipsis 的单行代码同时适用于 3D 和 2D 数组 -

foo[(Ellipsis,)+np.triu_indices(2)]

它同样适用于 2D 数组 -

out = foo[r,c] # foo as 2D input array

掩蔽方式

3D 数组案例

我们还可以为基于 masking 的方式使用掩码 -

foo[:,~np.tri(2,k=-1, dtype=bool)]

二维数组案例

foo[~np.tri(2,k=-1, dtype=bool)]

【讨论】：