Как получить индекс максимального элемента в массиве numpy вдоль одной оси

У меня есть двумерный массив NumPy. Я знаю, как получить максимальные значения по осям:

>>> a = array([[1,2,3],[4,3,1]])
>>> amax(a,axis=0)
array([4, 3, 3])

Как мне получить индексы максимальных элементов? Поэтому я хотел бы в качестве вывода array ([1,1,0])