y,i=lab.max(x)
returns a tensor y of the largest element in
each column of x, and a tensor i of their corresponding indices in x.
y,i=lab.max(x,2)
performs the max operation for each row and
y,i=lab.max(x,n)
performs the max operation over the dimension n.