a = np.array([[np.nan, 4], [2, 3]])
print(np.argmax(a))
0
print(np.nanargmax(a))
1
np.nanargmax(a, axis=0)
array([1, 0])
np.nanargmax(a, axis=1)
array([1, 1])
Finding index of maximum value in array with NumPy