Python 沿出现最大值的轴获取索引

我正在使用python和numpy进行一些数据分析。那么，假设我有以下部分：

import numpy as np
x = np.random.rand(2, 2, 2)

导致：

array([[[0.7213753 , 0.89782739],
        [0.10375189, 0.02501165]],

       [[0.732744  , 0.17957702],
       [0.85643144, 0.7516079 ]]])

array([[0.732744  , 0.89782739],
       [0.85643144, 0.7516079 ]])

现在，我可以沿着第一个轴找到最大值，如下所示：

np.max(x, axis=0)

导致：

array([[[0.7213753 , 0.89782739],
        [0.10375189, 0.02501165]],

       [[0.732744  , 0.17957702],
       [0.85643144, 0.7516079 ]]])

array([[0.732744  , 0.89782739],
       [0.85643144, 0.7516079 ]])

但是，我想做的不是最大值，而是沿着最大值出现的第一个轴的索引。因此，对于该数据，输出应为：

[[1, 0
  1, 1]]

我试过这样的方法：

np.where(x == np.max(x, axis=0))

但这并不能提供所需格式的输出。

>x.argmax（axis=0） 数组（[[1，0]， [1, 1]])