Rescaling alphas for Proxyless-Gradient

rescale.py

def rescale(a, a_new, i_op_0, i_op_1):
    i_op_0 = int(i_op_0)
    i_op_1 = int(i_op_1)
    old_p = F.softmax(a, dim=0)
    new_p = F.softmax(a_new, dim=0)
    old_sum = old_p[i_op_0] + old_p[i_op_1]
    new_sum = new_p[i_op_0] + new_p[i_op_1]
    ratio = old_sum / new_sum
    # rescaled probabilties such that sum is same as before
    p_r_0 = ratio * new_p[i_op_0]
    p_r_1 = ratio * new_p[i_op_1]
    new_sum_a = sum([th.exp(a[i]) for i in range(len(a)) if i not in [i_op_0, i_op_1]])

    new_a_0 = th.log((new_sum_a * (p_r_0 + ((p_r_0 * p_r_1)/(1-p_r_1)))) / (
        1 - p_r_0 - ((p_r_0*p_r_1) / (1-p_r_1))))
    new_a_1 = th.log((new_sum_a * (p_r_1 + ((p_r_1 * p_r_0)/(1-p_r_0)))) / (
        1 - p_r_1 - ((p_r_1*p_r_0) / (1-p_r_0))))

    a_new.data[i_op_0] = new_a_0
    a_new.data[i_op_1] = new_a_1