def _time_auto_unit(time_max):
    """
    Automatically set the unit for time axis according to the absolute maximum
    of the input data. This is a separate function for ease of testing and for
    use across different plots.

    Parameters
    ----------

    time_max : float
        Absolute maximum of the time data in seconds
    """

    if time_max == 0:
        unit = 's'
    elif time_max < 1e-3:
        unit = 'mus'
    elif time_max < 1:
        unit = 'ms'
    else:
        unit = 's'

    return unit