How to get the index of a maximum element in a NumPy array along one axis

>>> import numpy as np
>>> a = np.array([[1,2,3],[4,3,1]])
>>> i,j = np.unravel_index(a.argmax(), a.shape)
>>> a[i,j]
4

>>> a.argmax(axis=0)

array([1, 1, 0])