Returns index which holds maximum of array elements over a given axis.
x (Variable or N-dimensional array) – Array to find maximum elements.
axis (None or int) – Axis over which a max is performed.
The default (axis = None) is perform a max over all the dimensions
of the input array.