import torch


def factor_svd(original_layer, assert_equal=True):
    """
    Returns an SVD decomposition where we have multiplied the U and V vectors
    by sqrt(S) in order to distribute its weighting across the two evenly.

    If assert_equal is true, it will also verify that 99% of the weights of the
    recombined US^1/2 * S^1/2V^T matrix are within 1e-3 of the original weights.
    Those bounds are mostly arbitrary, just what I found empirically.
    """
    W = original_layer.state_dict()['weight']
    svd_ret = torch.svd(W)
    U,V,S, = svd_ret.U, svd_ret.V, svd_ret.S
    sqrt_lambda = torch.diag(S ** (1/2))
    ULhalf = torch.matmul(U, sqrt_lambda)
    LhalfVT = torch.matmul(sqrt_lambda, V.T)
    if assert_equal:
        rtol=1e-3
        target_ratio = 0.98 # should try for 0.98 but the finetuned weights that we load in only hit like 0.78 for some reason?
        ratio = float(torch.isclose(torch.matmul(ULhalf,LhalfVT), W, rtol=rtol).flatten().sum()) / W.flatten().size()[0]
        # assert ratio >= target_ratio
        print("ratio {} vs target_ratio {}".format(ratio, target_ratio))
    return ULhalf, LhalfVT