modified_uniform_quantization.py

def uniform_q(x, bits=32):
    if bits == 32:
        return x

    n = float(2 ** bits)
    out = np.clip(np.round(x * n - 0.5), 0, n-1) / (n-1)
    return out
    
def quantize_weights(x, bits=32):
    if bits == 32:
        return x

    # preprocess weight before quantization
    weight = np.tanh(x)
    weight = (weight / (2 * np.max(np.abs(weight)))) + 0.5 
    # quantize
    weight_q = 2 * uniform_q(weight, bits=bits) - 1 # w_q equation (1) <--- see paper

    return weight_q