import numpy as np
x = np.array([[5,2,1,3], [2,1,5]])
fun = lambda t: np.argmax(t)
np.array([fun(xi) for xi in x])