Returns indices of the maximum values along an axis

def argmax(attrs, inputs, proto_obj):
    """Returns indices of the maximum values along an axis"""
    axis = attrs.get('axis', 0)
    keepdims = attrs.get('keepdims', 1)
    argmax_op = symbol.argmax(inputs[0], axis=axis, keepdims=keepdims)
    # onnx argmax operator always expects int64 as output type
    cast_attrs = {'dtype': 'int64'}
    return 'cast', cast_attrs, argmax_op