def replace_labels(axes, cols, labels):
    """
    Sets the the labels of the graph in a proper format. For example removes 'mag' from the old labels

    :param axes: The axes object of the subplot

    :param cols: The names of the cols
    :type cols: list
    :param labels: The new labels
    :type labels: list, dict
    :return:
    """
    if type(labels) == list:
        axes.set_xlabel(labels[0])
        axes.set_ylabel(labels[1])
    elif type(labels) == dict:
        axes.set_xlabel(labels[cols[0]])
        axes.set_ylabel(labels[cols[1]])
    else:
        axes.set_xlabel(cols[0].replace('mag', ''))
        axes.set_ylabel(cols[1].replace('mag', ''))
