Function argmax

pub fn argmax<T: ArgReduceArgs>(x: &Tensor, args: T) -> Tensor