'Calculates the indices of the maximal element value along the specified
axis.  If multiple elements contain the maximum, only the indices of
the first are returned.
Arguments:
a (Tensor): the Tensor on which to perform the operation
axis (int, optional): the dimension along which to compute.
If set to None, we will take argmax over all
dimensions.  Defaults to 1
out (Tensor, optional): where the result will be stored. If out is
None, only the op-tree will be returned.
keepdims (bool, optional): Keep the axes being computed over in the
output (with size 1), instead of
collapsing.  Defaults to True.
Returns:
OpTreeNode: the resulting op-tree'

def argmax(self, a, axis=1, out=None, keepdims=True):

    return OpTreeNode.build('argmax', a, None, axis=axis, out=out)
