get the index along the axes where the maximum value happens

Question

I am using python and numpy for some data analysis. So, say I have the following segment:

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

resulting in:

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

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

Now, I can find the maximum along the first axes as:

np.max(x, axis=0)

resulting in:

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

However, what I want to do is instead of the maximum value, I want the index along the first axes where the maximum occurs. So for this data, the output should be:

[[1, 0
  1, 1]]

I tried something like:

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

but this does not give the output in the desired format.

Answer 1

>>> x.argmax(axis=0)
array([[1, 0],
       [1, 1]])