我有一个三维 numpy 数组。有什么方法可以最快地获取一个 3D 数组，该数组具有数组每个最终轴的最大项，而无需编写循环。（我稍后将使用具有相同语法的 CuPy，循环将带走 GPU 并行性和速度，这是这里最重要的因素。）
获取最大项目的索引很容易：
>>> arr = np.array(
[[[ 6, -2, -6, -5],
[ 1, 12, 3, 9],
[21, 7, 9, 8]],
[[15, 12, 20, 12],
[17, 15, 17, 23],
[22, 18, 27, 32]]])
>>> indexes = arr.argmax(axis=2, keepdims=True)
>>> indexes
array([[[0],
[1],
[0]],
[[2],
[3],
[3]]])
但是如何使用这些索引从 arr 中获取选定的值？我尝试过的所有方法要么产生错误（例如 arr[indexes]），要么产生错误的结果。我希望在这个例子中得到的是
array([[[6],
[12],
[21]],
[[20],
[23],
[32]]])