获取轴上 numpy.argmax 元素的索引

我有 N 维矩阵，其中包含具有 N 个参数的函数的值。每个参数都有离散数量的值。我需要在除一个参数之外的所有参数上最大化该函数，从而产生一个大小等于非最大化参数值的数量的一维向量。我还需要保存其他参数采用的值。

为此，我想迭代应用numpy.max在不同的轴上降低矩阵的维数以找到我需要的东西。最终的向量将仅取决于我遗漏的参数。

然而，我无法找到最终元素的原始索引（其中包含有关其他参数所取值的信息）。我想使用numpy.argmax在相同的方式numpy.max但我无法取回原始索引。

我正在尝试的一个例子是：

x = [[[1,2],[0,1]],[[3,4],[6,7]]]
args = np.argmax(x, 0)

这返回

[[1 1]
 [1 1]]

这意味着 argmax 正在选择原始矩阵中的元素 (2,1,4,7)。但如何获得他们的指数呢？我试过unravel_index， 使用args直接作为矩阵的索引x，一堆从 numpy 到 index 的函数都没有成功。

Using numpy.where不是一个解决方案，因为输入矩阵内部可能具有相同的值，所以我无法区分不同的原始值。

x.argmax(0)给出沿第 1 轴的最大值索引。使用np.indices生成另一个轴的索引。

x = np.array([[[1,2],[0,1]],[[3,4],[6,7]]])
x.argmax(0)
    array([[1, 1],
           [1, 1]])
a1, a2 = np.indices((2,2))
(x.argmax(0),a1,a2)
    (array([[1, 1],
            [1, 1]]),
     array([[0, 0],
            [1, 1]]),
     array([[0, 1],
            [0, 1]]))


x[x.argmax(0),a1,a2]
    array([[3, 4],
           [6, 7]])

x[a1,x.argmax(1),a2] 
    array([[1, 2],
           [6, 7]])

x[a1,a2,x.argmax(2)] 
    array([[2, 1],
           [4, 7]])

If x有其他维度，生成a1, and a2适当地。

官方文档并没有说太多如何使用argmax，但早期的 SO 线程已经讨论过它。我从这里得到了这个总体想法在多维数组上使用 numpy.argmax() https://stackoverflow.com/questions/5798364/using-numpy-argmax-on-multidimensional-arrays