Implement argmax

[mortont]

Having Nx.argmax implemented will help to reduce backend_transfer calls when recursively calling Ortex for applications that use greedy search, beam search, etc.