def map_noise_scale_to_time_step(alpha_infer, alpha):
    """
    Map an alpha_infer to an approx. time step in the alpha tensor.
        (Modified from Algorithm 3 Fast Sampling in DiffWave)

    Parameters:
        alpha_infer (float): noise scale at time `n` for inference
        alpha (tensor):      noise scales used in training, shape=(T, )
    Returns:
        t (float):           approximated time step in alpha tensor
    """

    if alpha_infer < alpha[-1]:
        return len(alpha) - 1
    if alpha_infer > alpha[0]:
        return 0
    for t in range(len(alpha) - 1):
        if alpha[t+1] <= alpha_infer <= alpha[t]:
            step_diff = alpha[t] - alpha_infer
            step_diff /= alpha[t] - alpha[t+1]
            return t + step_diff.item()
    return -1