python - 索引错误 : shape mismatch: indexing arrays could not be broadcast together with shapes

Question

我正在尝试从张量中提取某些数据，但出现了奇怪的错误。在这里，我将尝试生成错误:

a=np.random.randn(5, 10, 5, 5)
a[:, [1, 6], np.triu_indices(5, 0)[0], np.triu_indices(5, 0)[1]].shape

我收到这个错误

shape mismatch: indexing arrays could not be broadcast together with shapes

我什至没有做任何广播!都是切片的东西。

我想要什么？保持第零轴不变(获取所有内容)，从第一个轴获取 [1] 和 [6]，通过仅采用上三角元素将最后两个轴从 [5, 5] reshape 为 [15]。

Answer 1

我们需要将第二个轴索引数组扩展到 2D，以便它形成一个 外平面 反对 np.triu_indices。因此，它为我们提供了一个 2D 网格 mxn 数组，其中 m 是第二个轴索引数组的长度，n 是 np.triu_indices 的长度。所以，从本质上讲，整个解决方案将简化为这样的东西 -

r,c = np.triu_indices(5, 0)
out = a[:, np.array([1, 6])[:,None], r, c]

或者将该扩展版本作为列表输入，即 -

out = a[:, [[1],[6]], r, c]

我们还可以使用基于掩码的np.tri/np.triu，这在较大的数组上可能会更快，因为我们会跳过创建所有整数索引，就像这样-

mask = ~np.tri(5, k=-1, dtype=bool)
out = a[:, np.array([1, 6])[:,None], mask]