jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def zca_forward(self, x):
        batch_size, feature_dim = x.size()
        eps = 0.1
        f_cov = (torch.mm(x.transpose(0, 1), x) / (batch_size - 1)).float()  # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        # f_cov = torch.mm(x.transpose(0, 1), x).float()  # d * d
        U, S, V = torch.linalg.svd(0.9 * f_cov + 0.1 * eye)
        diag = torch.diag(1.0 / torch.sqrt(S+1e-5))
        rotate_mtx = torch.mm(torch.mm(U, diag), U.transpose(0, 1)).detach() # d * d
        return torch.mm(rotate_mtx, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def _svd_weighted_adapter(
            self,
            adapters,
            weights,
            new_rank,
            target,
            target_lora_A,
            target_lora_B,
            clamp=None,
            full_matrices=True,
            driver=None,
    ):
        valid_adapters = []
        valid_weights = []
        for adapter, weight in zip(adapters, weights):
            if adapter in target.lora_A or adapter in target.lora_embedding_A:
                valid_adapters.append(adapter)
                valid_weights.append(weight)

        # if no valid adapter, nothing to do
        if len(valid_adapters) == 0:
            raise ValueError("No matching LoRAs found. Please raise an issue on Github.")

        delta_weight = valid_weights[0] * target.get_delta_weight(valid_adapters[0])
        for adapter, weight in zip(valid_adapters[1:], valid_weights[1:]):
            delta_weight += weight * target.get_delta_weight(adapter)
        conv2d = isinstance(target, Conv2d)
        if conv2d:
            conv2d_1x1 = target.weight.size()[2:4] == (1, 1)
            if not conv2d_1x1:
                delta_weight = delta_weight.flatten(start_dim=1)
            else:
                delta_weight = delta_weight.squeeze()
        if hasattr(target, "fan_in_fan_out") and target.fan_in_fan_out:
            delta_weight = delta_weight.T

        # based on https://github.com/kohya-ss/sd-scripts/blob/main/networks/svd_merge_lora.py#L114-L131
        U, S, Vh = torch.svd(delta_weight, some=full_matrices)
        U = U[:, :new_rank]
        S = S[:new_rank]
        U = U @ torch.diag(S)
        Vh = Vh[:new_rank, :]
        if clamp is not None:
            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp)
            low_val = -hi_val
            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)
        if conv2d:
            U = U.reshape(target_lora_B.data.shape)
            Vh = Vh.reshape(target_lora_A.data.shape)
        return Vh, U
//  [python]
def _svd_weighted_adapter(
            self,
            adapters,
            weights,
            new_rank,
            target,
            target_lora_A,
            target_lora_B,
            clamp=None,
            full_matrices=True,
            driver=None,
    ):
        valid_adapters = []
        valid_weights = []
        for adapter, weight in zip(adapters, weights):
            if adapter in target.lora_A or adapter in target.lora_embedding_A:
                valid_adapters.append(adapter)
                valid_weights.append(weight)

        # if no valid adapter, nothing to do
        if len(valid_adapters) == 0:
            raise ValueError("No matching LoRAs found. Please raise an issue on Github.")

        delta_weight = valid_weights[0] * target.get_delta_weight(valid_adapters[0])
        for adapter, weight in zip(valid_adapters[1:], valid_weights[1:]):
            delta_weight += weight * target.get_delta_weight(adapter)
        conv2d = isinstance(target, Conv2d)
        if conv2d:
            conv2d_1x1 = target.weight.size()[2:4] == (1, 1)
            if not conv2d_1x1:
                delta_weight = delta_weight.flatten(start_dim=1)
            else:
                delta_weight = delta_weight.squeeze()
        if hasattr(target, "fan_in_fan_out") and target.fan_in_fan_out:
            delta_weight = delta_weight.T

        # based on https://github.com/kohya-ss/sd-scripts/blob/main/networks/svd_merge_lora.py#L114-L131
        U, S, Vh = torch.linalg.svd(delta_weight, full_matrices=full_matrices, driver=driver)
        U = U[:, :new_rank]
        S = S[:new_rank]
        U = U @ torch.diag(S)
        Vh = Vh[:new_rank, :]
        if clamp is not None:
            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp)
            low_val = -hi_val
            U = U.clamp(low_val, hi_val)
            Vh = Vh.clamp(low_val, hi_val)
        if conv2d:
            U = U.reshape(target_lora_B.data.shape)
            Vh = Vh.reshape(target_lora_A.data.shape)
        return Vh, U
//  [python]
def normalizeSE3_inplace(T):
    R = T[..., :3, :3]
    U, S, Vh = torch.svd(R)
    T[..., :3, :3] = torch.matmul(U, Vh)
//  [python]
def normalizeSE3_inplace(T):
    R = T[..., :3, :3]
    U, S, Vh = torch.linalg.svd(R)
    T[..., :3, :3] = torch.matmul(U, Vh)
//  [python]
def get_mf_mode(params, num_samples=0):
    """ Get the mode of matrix fisher distribution.
    U, S, Vh = F
    R_est = U @ [[1,0,0],[0,1,0],[0,0,|UV|]] @ Vh, where Vh = V.T
    
    Args:
    - params: (B, 3, 3), the parameter of matrix-Fisher distribution.
    - num_samples: scalar.

    Returns:
    - mode (B, 3, 3) or samples (B, S+1, 3, 3), where the first is the mode.

    ATTN 1: If the condition number of F is large, the solution of svd may be unstable.
    ATTN 2: The api of torch.linalg.svd (U, S, Vh) is different from torch.svd (U, S, V).
    ATTN 3: Batch svd is somewhat different from single svd.
    """

    batch_size = params.shape[0]
    
    diag_mat = 1e-5 * torch.eye(3, device=params.device).unsqueeze(0).expand(batch_size, -1, -1).clone()
    U, S, Vh = torch.svd(params + diag_mat)

    with torch.no_grad():
        det_u_v = torch.linalg.det(torch.bmm(U, Vh.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(batch_size, -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v

    mode = torch.bmm(torch.bmm(U, det_modify_mat), Vh)

    if num_samples > 0:
        V = Vh.transpose(-2, -1)
        # samples : (B, S, 3, 3)
        samples = mf_sampling_torch(U, S, V, num_samples, b=1.5, oversampling_ratio=8, sample_on_cpu=False)

        return torch.concat((mode.unsqueeze(1), samples), dim=1)

    return mode
//  [python]
def get_mf_mode(params, num_samples=0):
    """ Get the mode of matrix fisher distribution.
    U, S, Vh = F
    R_est = U @ [[1,0,0],[0,1,0],[0,0,|UV|]] @ Vh, where Vh = V.T
    
    Args:
    - params: (B, 3, 3), the parameter of matrix-Fisher distribution.
    - num_samples: scalar.

    Returns:
    - mode (B, 3, 3) or samples (B, S+1, 3, 3), where the first is the mode.

    ATTN 1: If the condition number of F is large, the solution of svd may be unstable.
    ATTN 2: The api of torch.linalg.svd (U, S, Vh) is different from torch.svd (U, S, V).
    ATTN 3: Batch svd is somewhat different from single svd.
    """

    batch_size = params.shape[0]
    
    diag_mat = 1e-5 * torch.eye(3, device=params.device).unsqueeze(0).expand(batch_size, -1, -1).clone()
    U, S, Vh = torch.linalg.svd(params + diag_mat, full_matrices=True)

    with torch.no_grad():
        det_u_v = torch.linalg.det(torch.bmm(U, Vh.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(batch_size, -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v

    mode = torch.bmm(torch.bmm(U, det_modify_mat), Vh)

    if num_samples > 0:
        V = Vh.transpose(-2, -1)
        # samples : (B, S, 3, 3)
        samples = mf_sampling_torch(U, S, V, num_samples, b=1.5, oversampling_ratio=8, sample_on_cpu=False)

        return torch.concat((mode.unsqueeze(1), samples), dim=1)

    return mode
//  [python]
def SVDShrink_Fast_GPU(self, X, tau):
        D = torch.fft.fft(X)
        U_bar = torch.empty(X.shape[0], X.shape[1], 0).to(device)
        S_bar = torch.empty(X.shape[0], X.shape[1], 0).to(device)
        V_bar = torch.empty(X.shape[0], X.shape[1], 0).to(device)
        W_bar = torch.empty((X.shape[0], X.shape[1], 0)).to(device)
        l = X.shape[2]
        for i in range(l):
            if i < math.ceil((l + 1) / 2):
                U, S, V = torch.svd(D[:, :, i])
                S = self.SoftShrink_GPU(S, tau)
                S = torch.diag(S).to(torch.complex64)
                w = self.tensor_mul(self.tensor_mul(U, S), V)
                # w = torch.dot(torch.dot(U,S),V)
                U_bar = torch.cat((U_bar, U.reshape(X.shape[0], X.shape[1], 1)), 2)
                S_bar = torch.cat((S_bar, S.reshape(X.shape[0], X.shape[1], 1)), 2)
                V_bar = torch.cat((V_bar, V.reshape(X.shape[0], X.shape[1], 1)), 2)
                W_bar = torch.cat((W_bar, w.reshape(X.shape[0], X.shape[1], 1)), 2)
            elif i >= math.ceil((l + 1) / 2):
                U = torch.conj(U_bar[:, :, l - i])
                S = S_bar[:, :, l - i]
                V = torch.conj(V_bar[:, :, l - i])
                w = self.tensor_mul(self.tensor_mul(U, S), V)
                W_bar = torch.cat((W_bar, w.reshape(X.shape[0], X.shape[1], 1)), 2)
        return torch.fft.ifft(W_bar).real
//  [python]
def SVDShrink_Fast_GPU(self, X, tau):
        D = torch.fft.fft(X)
        U_bar = torch.empty(X.shape[0], X.shape[1], 0).to(device)
        S_bar = torch.empty(X.shape[0], X.shape[1], 0).to(device)
        V_bar = torch.empty(X.shape[0], X.shape[1], 0).to(device)
        W_bar = torch.empty((X.shape[0], X.shape[1], 0)).to(device)
        l = X.shape[2]
        for i in range(l):
            if i < math.ceil((l + 1) / 2):
                U, S, V = torch.linalg.svd(D[:, :, i])
                S = self.SoftShrink_GPU(S, tau)
                S = torch.diag(S).to(torch.complex64)
                w = self.tensor_mul(self.tensor_mul(U, S), V)
                # w = torch.dot(torch.dot(U,S),V)
                U_bar = torch.cat((U_bar, U.reshape(X.shape[0], X.shape[1], 1)), 2)
                S_bar = torch.cat((S_bar, S.reshape(X.shape[0], X.shape[1], 1)), 2)
                V_bar = torch.cat((V_bar, V.reshape(X.shape[0], X.shape[1], 1)), 2)
                W_bar = torch.cat((W_bar, w.reshape(X.shape[0], X.shape[1], 1)), 2)
            elif i >= math.ceil((l + 1) / 2):
                U = torch.conj(U_bar[:, :, l - i])
                S = S_bar[:, :, l - i]
                V = torch.conj(V_bar[:, :, l - i])
                w = self.tensor_mul(self.tensor_mul(U, S), V)
                W_bar = torch.cat((W_bar, w.reshape(X.shape[0], X.shape[1], 1)), 2)
        return torch.fft.ifft(W_bar).real
//  [python]
def low_rank_approximate(weight, rank, clamp_quantile=0.99):
    if len(weight.shape)==4: # conv
        weight=weight.flatten(1)
        out_ch, in_ch, k1, k2 = weight.shape

    U, S, Vh = torch.svd(weight)
    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, clamp_quantile)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if len(weight.shape) == 4:
        # U is (out_channels, rank) with 1x1 conv.
        U = U.reshape(U.shape[0], U.shape[1], 1, 1)
        # V is (rank, in_channels * kernel_size1 * kernel_size2)
        Vh = Vh.reshape(Vh.shape[0], in_ch, k1, k2)
    return U, Vh
//  [python]
def low_rank_approximate(weight, rank, clamp_quantile=0.99):
    if len(weight.shape)==4: # conv
        weight=weight.flatten(1)
        out_ch, in_ch, k1, k2 = weight.shape

    U, S, Vh = torch.linalg.svd(weight)
    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, clamp_quantile)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if len(weight.shape) == 4:
        # U is (out_channels, rank) with 1x1 conv.
        U = U.reshape(U.shape[0], U.shape[1], 1, 1)
        # V is (rank, in_channels * kernel_size1 * kernel_size2)
        Vh = Vh.reshape(Vh.shape[0], in_ch, k1, k2)
    return U, Vh
//  [python]
def _svd_initialize(Wa: torch.Tensor, Wb: Optional[torch.Tensor], k: int) ->Tuple[torch.Tensor, torch.Tensor]:
    if Wb is None:
        W = Wa
    else:
        W = Wa @ Wb
    U, s, Vt = torch.svd(W)
    Wa = U[:, :k]
    Wb = Vt[:k]
    s_sqrt = torch.sqrt(s[:k])
    Wa *= s_sqrt
    Wb *= s_sqrt.reshape(-1, 1)
    return Wa, Wb
//  [python]
def _svd_initialize(Wa: torch.Tensor, Wb: Optional[torch.Tensor], k: int) ->Tuple[torch.Tensor, torch.Tensor]:
    if Wb is None:
        W = Wa
    else:
        W = Wa @ Wb
    U, s, Vt = torch.linalg.svd(W, full_matrices=False)
    Wa = U[:, :k]
    Wb = Vt[:k]
    s_sqrt = torch.sqrt(s[:k])
    Wa *= s_sqrt
    Wb *= s_sqrt.reshape(-1, 1)
    return Wa, Wb
//  [python]
def nrqm(
    img: Tensor,
    linear_param,
    rf_param,
) -> Tensor:
    """Calculate NRQM
    Args:
        img (Tensor): Input image.
        linear_param (np.array): (4, 1) linear regression params
        rf_param: params of 3 random forest for 3 kinds of features
    """
    assert (
        img.ndim == 4
    ), "Input image must be a gray or Y (of YCbCr) image with shape (b, c, h, w)."

    # crop image
    b, c, h, w = img.shape
    img_pyr = get_guass_pyramid(img.float() / 255.0)

    # DCT features
    f1 = []
    for im in img_pyr:
        f1.append(block_dct(im))
    f1 = torch.cat(f1, dim=1)

    # gsm features
    f2 = global_gsm(img)

    # svd features
    f3 = []
    for im in img_pyr:
        col = im2col(im, 5, "distinct")
        _, s, _ = torch.svd(col)
        f3.append(s)
    f3 = torch.cat(f3, dim=1)

    # Random forest regression. Currently not differentiable and only support CPU
    preds = torch.ones(b, 1)
    for feat, rf in zip([f1, f2, f3], rf_param):
        tmp_pred = random_forest_regression(feat, *rf)
        preds = torch.cat((preds, tmp_pred), dim=1)
    quality = preds @ torch.Tensor(linear_param)

    return quality.squeeze()
//  [python]
def nrqm(
    img: Tensor,
    linear_param,
    rf_param,
) -> Tensor:
    """Calculate NRQM
    Args:
        img (Tensor): Input image.
        linear_param (np.array): (4, 1) linear regression params
        rf_param: params of 3 random forest for 3 kinds of features
    """
    assert (
        img.ndim == 4
    ), "Input image must be a gray or Y (of YCbCr) image with shape (b, c, h, w)."

    # crop image
    b, c, h, w = img.shape
    img_pyr = get_guass_pyramid(img.float() / 255.0)

    # DCT features
    f1 = []
    for im in img_pyr:
        f1.append(block_dct(im))
    f1 = torch.cat(f1, dim=1)

    # gsm features
    f2 = global_gsm(img)

    # svd features
    f3 = []
    for im in img_pyr:
        col = im2col(im, 5, "distinct")
        _, s, _ = torch.linalg.svd(col, full_matrices=False)
        f3.append(s)
    f3 = torch.cat(f3, dim=1)

    # Random forest regression. Currently not differentiable and only support CPU
    preds = torch.ones(b, 1)
    for feat, rf in zip([f1, f2, f3], rf_param):
        tmp_pred = random_forest_regression(feat, *rf)
        preds = torch.cat((preds, tmp_pred), dim=1)
    quality = preds @ torch.Tensor(linear_param)

    return quality.squeeze()
//  [python]
def eigh(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:  # pragma: no cover
    """
    Like torch.linalg.eigh, assuming the argument is a symmetric real matrix.
    """
    #if hasattr(torch, "linalg") and hasattr(torch.linalg, "eigh"):
    #    return torch.linalg.eigh(A)
    #return torch.symeigh(A)


    #S, U = torch.linalg.eigh(A)
    #print (S.shape)
    #print (U.shape)
    #print (S[0], U[0])
    U, S, _ = torch.svd(A)
    #print (S.shape)
    #print (U.shape)
    S = torch.flip(S, (-1, ))
    U = torch.flip(U, (-1, ))

    #print (S[0], U[0])
   
    return S, U
//  [python]
def eigh(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:  # pragma: no cover
    """
    Like torch.linalg.eigh, assuming the argument is a symmetric real matrix.
    """
    #if hasattr(torch, "linalg") and hasattr(torch.linalg, "eigh"):
    #    return torch.linalg.eigh(A)
    #return torch.symeigh(A)


    #S, U = torch.linalg.eigh(A)
    #print (S.shape)
    #print (U.shape)
    #print (S[0], U[0])
    U, S, _ = torch.linalg.svd(A)
    #print (S.shape)
    #print (U.shape)
    S = torch.flip(S, (-1, ))
    U = torch.flip(U, (-1, ))

    #print (S[0], U[0])
   
    return S, U
//  [python]
def _svd(w, gain=1):
    rows = w.size(0)
    cols = w.numel() // rows
    flattened = w.new(rows, cols).normal_(0, 1)
    if rows > cols:
        flattened.t_()
        u, _, v = torch.svd(flattened)
        u.t_()
        v.t_()
    else:
        u, _, v = torch.svd(flattened)
    q = u if u.shape == w.shape else v
    q = q.reshape(w.shape)

    with torch.no_grad():
        w.view_as(q).copy_(q)
        w.mul_(gain)
    return w
//  [python]
def _svd(w, gain=1):
    rows = w.size(0)
    cols = w.numel() // rows
    flattened = w.new(rows, cols).normal_(0, 1)
    if rows > cols:
        flattened.t_()
        u, _, v = torch.linalg.svd(flattened, full_matrices=False)
        u.t_()
        v.t_()
    else:
        u, _, v = torch.linalg.svd(flattened, full_matrices=False)
    q = u if u.shape == w.shape else v
    q = q.reshape(w.shape)

    with torch.no_grad():
        w.view_as(q).copy_(q)
        w.mul_(gain)
    return w
//  [python]
def forward(self, M, mask):
        loss = 0
        batch_size, C, H, W = M.shape
        avg = mean_image(M)
        #M = M - avg
        M = M * mask
        M = M.reshape((batch_size, C, H * W))
        M = M.reshape((batch_size, C * H * W))
        M1 = torch.transpose(M, 0, 1)#.detach()
        u, s, vh = torch.svd(M1)
        s[1:] = 0
        loss = torch.dist(torch.transpose(M, 0, 1), u @ torch.diag(s) @ vh)/torch.sum(mask)
        return loss*1e2
//  [python]
def forward(self, M, mask):
        loss = 0
        batch_size, C, H, W = M.shape
        avg = mean_image(M)
        #M = M - avg
        M = M * mask
        M = M.reshape((batch_size, C, H * W))
        M = M.reshape((batch_size, C * H * W))
        M1 = torch.transpose(M, 0, 1)#.detach()
        u, s, vh = torch.linalg.svd(M1, full_matrices=False)
        s[1:] = 0
        loss = torch.dist(torch.transpose(M, 0, 1), u @ torch.diag(s) @ vh)/torch.sum(mask)
        return loss*1e2
//  [python]
def forward(self, M, mask):
        loss = 0
        batch_size, C, H, W = M.shape
        avg = mean_image(M)
        #M = M - avg
        M = M * mask
        M = M.reshape((batch_size, C, H * W))
        M = M.reshape((batch_size, C * H * W))
        M1 = torch.transpose(M, 0, 1).detach()
        u, s, vh = torch.svd(M1)
        s[1:] = 0
        loss = torch.dist(torch.transpose(M, 0, 1), u @ torch.diag(s) @ vh)/torch.sum(mask)
        return loss*1e2
//  [python]
def forward(self, M, mask):
        loss = 0
        batch_size, C, H, W = M.shape
        avg = mean_image(M)
        #M = M - avg
        M = M * mask
        M = M.reshape((batch_size, C, H * W))
        M = M.reshape((batch_size, C * H * W))
        M1 = torch.transpose(M, 0, 1).detach()
        u, s, vh = torch.linalg.svd(M1, full_matrices=False)
        s[1:] = 0
        loss = torch.dist(torch.transpose(M, 0, 1), u @ torch.diag(s) @ vh)/torch.sum(mask)
        return loss*1e2
//  [python]
def forward(self, I, S, mask):
        loss = 0
        M0 = I-S
        M0 = M0*mask
        M0 = torch.clip(M0, 0.00001, 1.0)
        M = M0[:,:2,:,:]
        Msum = torch.sum(M0,1,keepdim=True)#.detach()
        Msum = torch.tile(Msum, (1, 2,1,1))
        M = M/Msum
        batch_size, C, H, W = M.shape
        M = M * mask
        
        M = M.reshape((batch_size, C, H * W))
        M = M.reshape((batch_size, C * H * W))
        M1 = torch.transpose(M, 0, 1).detach()
        u, s, vh = torch.svd(M1)
        s[1:] = 0
        loss = torch.dist(torch.transpose(M, 0, 1), u @ torch.diag(s) @ vh)/torch.sum(mask)
        return loss*1e2
//  [python]
def forward(self, I, S, mask):
        loss = 0
        M0 = I-S
        M0 = M0*mask
        M0 = torch.clip(M0, 0.00001, 1.0)
        M = M0[:,:2,:,:]
        Msum = torch.sum(M0,1,keepdim=True)#.detach()
        Msum = torch.tile(Msum, (1, 2,1,1))
        M = M/Msum
        batch_size, C, H, W = M.shape
        M = M * mask
        
        M = M.reshape((batch_size, C, H * W))
        M = M.reshape((batch_size, C * H * W))
        M1 = torch.transpose(M, 0, 1).detach()
        u, s, vh = torch.linalg.svd(M1, full_matrices=False)
        s[1:] = 0
        loss = torch.dist(torch.transpose(M, 0, 1), u @ torch.diag(s) @ vh)/torch.sum(mask)
        return loss*1e2
//  [python]
def svdtf(source, target):
    assert source.size(-2) == target.size(-2), {"The number of points N has to be the same for both point clouds."}
    ctnsource = source.mean(dim=-2, keepdim=True)
    ctntarget = target.mean(dim=-2, keepdim=True)
    source = source - ctnsource
    target = target - ctntarget
    M = torch.einsum('...Na, ...Nb -> ...ab', target, source)
    U, S, Vh = torch.svd(M)
    R = U @ torch.diag_embed(torch.sign(S)) @ Vh
    mask = (R.det() + 1).abs() < 1e-6
    R[mask] = - R[mask]
    t = ctntarget.mT - R @ ctnsource.mT
    T = torch.cat((R, t), dim=-1)
    return mat2SE3(T, check=False)
//  [python]
def svdtf(source, target):
    assert source.size(-2) == target.size(-2), {"The number of points N has to be the same for both point clouds."}
    ctnsource = source.mean(dim=-2, keepdim=True)
    ctntarget = target.mean(dim=-2, keepdim=True)
    source = source - ctnsource
    target = target - ctntarget
    M = torch.einsum('...Na, ...Nb -> ...ab', target, source)
    U, S, Vh = torch.linalg.svd(M)
    R = U @ Vh
    mask = (R.det() + 1).abs() < 1e-6
    R[mask] = - R[mask]
    t = ctntarget.mT - R @ ctnsource.mT
    T = torch.cat((R, t), dim=-1)
    return mat2SE3(T, check=False)
//  [python]
def build_usv_from_weight(self, weight: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        ### differentiable feature is gauranteed
        U, S, V = weight.data.svd(some=False)
        V = V.transpose(-2, -1).contiguous()
        # U, S, V = torch.linalg.svd(weight.data, full_matrices=True, driver="gesvd") # must use QR decomposition
        self.U.data.copy_(U)
        self.S.data.copy_(S)
        self.V.data.copy_(V)
        return U, S, V
//  [python]
def build_usv_from_weight(self, weight: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        ### differentiable feature is gauranteed
        # U, S, V = weight.data.svd(some=False)
        # V = V.transpose(-2, -1).contiguous()
        U, S, V = torch.linalg.svd(weight.data, full_matrices=True, driver="gesvd") # must use QR decomposition
        self.U.data.copy_(U)
        self.S.data.copy_(S)
        self.V.data.copy_(V)
        return U, S, V
//  [python]
def reset_parameters(self) -> None:
        if self.mode == "weight":
            init.kaiming_normal_(self.weight.data)
        elif self.mode == "usv":
            W = init.kaiming_normal_(
                torch.empty(self.out_channels, self.in_channels_flat, dtype=self.U.dtype, device=self.device)
            )
            U, S, V = torch.svd(W, some=False)
            V = V.transpose(-2, -1)
            self.U.data.copy_(U)
            self.V.data.copy_(V)
            self.S.data.copy_(torch.ones(S.shape[0], dtype=self.U.dtype, device=self.device))
        elif self.mode == "phase":
            W = init.kaiming_normal_(
                torch.empty(self.out_channels, self.in_channels_flat, dtype=self.U.dtype, device=self.device)
            )
            U, S, V = torch.svd(W, some=False)
            V = V.transpose(-2, -1)
            delta_list, phi_mat = self.decomposer.decompose(U)
            self.delta_list_U.data.copy_(delta_list)
            self.phase_U.data.copy_(self.decomposer.m2v(phi_mat))
            delta_list, phi_mat = self.decomposer.decompose(V)
            self.delta_list_V.data.copy_(delta_list)
            self.phase_V.data.copy_(self.decomposer.m2v(phi_mat))
            self.S_scale.data.copy_(S.abs().max(dim=-1, keepdim=True)[0])
            self.phase_S.data.copy_(S.div(self.S_scale.data).acos())

        elif self.mode == "voltage":
            raise NotImplementedError
        else:
            raise NotImplementedError

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def reset_parameters(self) -> None:
        if self.mode == "weight":
            init.kaiming_normal_(self.weight.data)
        elif self.mode == "usv":
            W = init.kaiming_normal_(
                torch.empty(self.out_channels, self.in_channels_flat, dtype=self.U.dtype, device=self.device)
            )
            # U, S, V = torch.svd(W, some=False)
            # V = V.transpose(-2, -1)
            U, S, V = torch.linalg.svd(W, full_matrices=True, driver="gesvd") # must use QR decomposition
            self.U.data.copy_(U)
            self.V.data.copy_(V)
            self.S.data.copy_(torch.ones(S.shape[0], dtype=self.U.dtype, device=self.device))
        elif self.mode == "phase":
            W = init.kaiming_normal_(
                torch.empty(self.out_channels, self.in_channels_flat, dtype=self.U.dtype, device=self.device)
            )
            # U, S, V = torch.svd(W, some=False)
            # V = V.transpose(-2, -1)
            U, S, V = torch.linalg.svd(W, full_matrices=True, driver="gesvd") # must use QR decomposition
            delta_list, phi_mat = self.decomposer.decompose(U)
            self.delta_list_U.data.copy_(delta_list)
            self.phase_U.data.copy_(self.decomposer.m2v(phi_mat))
            delta_list, phi_mat = self.decomposer.decompose(V)
            self.delta_list_V.data.copy_(delta_list)
            self.phase_V.data.copy_(self.decomposer.m2v(phi_mat))
            self.S_scale.data.copy_(S.abs().max(dim=-1, keepdim=True)[0])
            self.phase_S.data.copy_(S.div(self.S_scale.data).acos())

        elif self.mode == "voltage":
            raise NotImplementedError
        else:
            raise NotImplementedError

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def decompose_E(model):
    try:
        u, s, vT = torch.svd(model)

    except Exception as e:
        print(e)
        model = torch.eye(3, device=model.device)
        u, s, vT = torch.svd(model)

    try:
        if torch.sum(torch.isnan(u)) >0:
            print("wrong")
    except Exception as e:
        print(e)

    w = torch.tensor([[0, -1, 0],
                      [1, 0, 0],
                      [0, 0, 1]], dtype=u.dtype, device=u.device)
    z = torch.tensor([[0, -1, 0],
                      [1, 0, 0],
                      [0, 0, 0]], dtype=u.dtype, device=u.device)


    u_ = u * (-1.0) if torch.det(u) < 0 else u

    vT_ = vT * (-1.0) if torch.det(vT) < 0 else vT

    R1 = u_ @ w @ vT_

    R2 = u_ @ w.transpose(0, 1) @ vT_

    t = u[:, -1]#real

    return R1, R2, t.unsqueeze(1)
//  [python]
def decompose_E(model):
    try:
        u, s, vT = torch.linalg.svd(model)

    except Exception as e:
        print(e)
        model = torch.eye(3, device=model.device)
        u, s, vT = torch.linalg.svd(model)

    try:
        if torch.sum(torch.isnan(u)) >0:
            print("wrong")
    except Exception as e:
        print(e)

    w = torch.tensor([[0, -1, 0],
                      [1, 0, 0],
                      [0, 0, 1]], dtype=u.dtype, device=u.device)
    z = torch.tensor([[0, -1, 0],
                      [1, 0, 0],
                      [0, 0, 0]], dtype=u.dtype, device=u.device)


    u_ = u * (-1.0) if torch.det(u) < 0 else u

    vT_ = vT * (-1.0) if torch.det(vT) < 0 else vT

    R1 = u_ @ w @ vT_

    R2 = u_ @ w.transpose(0, 1) @ vT_

    t = u[:, -1]#real

    return R1, R2, t.unsqueeze(1)
//  [python]
def reset_parameters(self) -> None:
        if self.mode == "weight":
            init.kaiming_normal_(self.weight.data)
        elif self.mode == "usv":
            W = init.kaiming_normal_(
                torch.empty(
                    self.grid_dim_y,
                    self.grid_dim_x,
                    self.miniblock,
                    self.miniblock,
                    dtype=self.U.dtype,
                    device=self.device,
                )
            )
            U, S, V = torch.svd(W, some=False)
            V = V.transpose(-2, -1)
            self.U.data.copy_(U)
            self.V.data.copy_(V)
            self.S.data.copy_(torch.ones_like(S, device=self.device))
        elif self.mode == "phase":
            W = init.kaiming_normal_(
                torch.empty(
                    self.grid_dim_y,
                    self.grid_dim_x,
                    self.miniblock,
                    self.miniblock,
                    dtype=self.U.dtype,
                    device=self.device,
                )
            )
            U, S, V = torch.svd(W, some=False)
            V = V.transpose(-2, -1)
            delta_list, phi_mat = self.decomposer.decompose(U)
            self.delta_list_U.data.copy_(delta_list)
            self.phase_U.data.copy_(self.decomposer.m2v(phi_mat))
            delta_list, phi_mat = self.decomposer.decompose(V)
            self.delta_list_V.data.copy_(delta_list)
            self.phase_V.data.copy_(self.decomposer.m2v(phi_mat))
            self.S_scale.data.copy_(S.abs().max(dim=-1, keepdim=True)[0])
            self.phase_S.data.copy_(S.div(self.S_scale.data).acos())

        elif self.mode == "voltage":
            raise NotImplementedError
        else:
            raise NotImplementedError

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def reset_parameters(self) -> None:
        if self.mode == "weight":
            init.kaiming_normal_(self.weight.data)
        elif self.mode == "usv":
            W = init.kaiming_normal_(
                torch.empty(
                    self.grid_dim_y,
                    self.grid_dim_x,
                    self.miniblock,
                    self.miniblock,
                    dtype=self.U.dtype,
                    device=self.device,
                )
            )
            # U, S, V = torch.svd(W, some=False)
            # V = V.transpose(-2, -1)
            U, S, V = torch.linalg.svd(W, full_matrices=True, driver="gesvd") # must use QR decomposition
            self.U.data.copy_(U)
            self.V.data.copy_(V)
            self.S.data.copy_(torch.ones_like(S, device=self.device))
        elif self.mode == "phase":
            W = init.kaiming_normal_(
                torch.empty(
                    self.grid_dim_y,
                    self.grid_dim_x,
                    self.miniblock,
                    self.miniblock,
                    dtype=self.U.dtype,
                    device=self.device,
                )
            )
            # U, S, V = torch.svd(W, some=False)
            # V = V.transpose(-2, -1)
            U, S, V = torch.linalg.svd(W, full_matrices=True, driver="gesvd") # must use QR decomposition
            delta_list, phi_mat = self.decomposer.decompose(U)
            self.delta_list_U.data.copy_(delta_list)
            self.phase_U.data.copy_(self.decomposer.m2v(phi_mat))
            delta_list, phi_mat = self.decomposer.decompose(V)
            self.delta_list_V.data.copy_(delta_list)
            self.phase_V.data.copy_(self.decomposer.m2v(phi_mat))
            self.S_scale.data.copy_(S.abs().max(dim=-1, keepdim=True)[0])
            self.phase_S.data.copy_(S.div(self.S_scale.data).acos())

        elif self.mode == "voltage":
            raise NotImplementedError
        else:
            raise NotImplementedError

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def estimate_minimal_model(self, pts, weights=None):  # x1 y1 x2 y2
        """
            using 5 points to estimate Essential matrix.
        """
        try:
            pts.shape[1] == self.sample_size

        except:
            print("This is not a minimal sample.")

        batch_size, num, _ = pts.shape
        dev = pts.device
        pts1 = pts[:, :, 0:2]
        pts2 = pts[:, :, 2:4]

        # get the points
        x1, y1 = pts1[:, :, 0], pts1[:, :, 1]
        x2, y2 = pts2[:, :, 0], pts2[:, :, 1]

        # Step1: construct the A matrix, A F = 0.
        # 5 equations for 9 variables, A is 5x9 matrix containing epipolar constraints
        # Essential matrix is a linear combination of the 4 vectors spanning the null space of A
        a_59 = torch.ones(x1.shape, device=dev)
        A_s = torch.stack(
            (torch.mul(x1, x2), torch.mul(x1, y2), x1,
             torch.mul(y1, x2), torch.mul(y1, y2), y1,
             x2, y2, a_59), dim=-1)

        U, S, V = torch.svd(A_s)#, full_matrices=False) #A_s # eigenvalues in increasing order

        null_space = V[:, -4:, :].transpose(-1, -2).float()  # the last four rows

        # use the 4 eigenvectors according to the 4 smallest singular values,
        # E is calculated from 4 basis, E = cx*X + cy*Y + cz*Z + cw*W, up to common scale = 1
        null_space_mat = null_space.reshape(null_space.shape[0], 3, 3, null_space.shape[-1]).transpose(1,
                                                                                                       2)  # X, Y, Z, W

        # Step2: expansion of the constraints:
        # determinant constraint det(E) = 0,
        # trace constraint 2EE^TE - trace(EE^T)E = 0
        constraint_mat = self.get_constraint_mat(null_space_mat, batch_size)

        # Step 3: Eliminate part of the matrix to isolate polynomials in z.
        # solve AX=b
        b = constraint_mat[:, :, 10:]
        eliminated_mat = torch.linalg.solve(constraint_mat[:, :, :10], b)

        action_mat = torch.zeros((batch_size, 10, 10), device=dev)
        action_mat[:, 0:3] = eliminated_mat[:, 0:3]
        action_mat[:, 3] = eliminated_mat[:, 4]
        action_mat[:, 4] = eliminated_mat[:, 5]
        action_mat[:, 5] = eliminated_mat[:, 7]
        action_mat[:, 6, 0] = -torch.ones_like(action_mat[:, 6, 0])
        action_mat[:, 7, 1] = -torch.ones_like(action_mat[:, 7, 1])
        action_mat[:, 8, 3] = -torch.ones_like(action_mat[:, 8, 3])
        action_mat[:, 9, 6] = -torch.ones_like(action_mat[:, 9, 6])
        ee, vv = torch.linalg.eig(action_mat)#, eigenvectors=True)# svd

        # put the cx, cy, cz back to get a valid essential matrix
        E_models = null_space.matmul(vv.real[:, -4:])#torch.stack(#.real
        real_mask = (vv.imag[:, -4:].abs() < 1e-3).all(dim=1).flatten()
        E_models = E_models.transpose(-1, -2).reshape(-1, 3, 3).transpose(-1, -2)

        return E_models[real_mask]
//  [python]
def estimate_minimal_model(self, pts, weights=None):  # x1 y1 x2 y2
        """
            using 5 points to estimate Essential matrix.
        """
        try:
            pts.shape[1] == self.sample_size

        except:
            print("This is not a minimal sample.")

        batch_size, num, _ = pts.shape
        dev = pts.device
        pts1 = pts[:, :, 0:2]
        pts2 = pts[:, :, 2:4]

        # get the points
        x1, y1 = pts1[:, :, 0], pts1[:, :, 1]
        x2, y2 = pts2[:, :, 0], pts2[:, :, 1]

        # Step1: construct the A matrix, A F = 0.
        # 5 equations for 9 variables, A is 5x9 matrix containing epipolar constraints
        # Essential matrix is a linear combination of the 4 vectors spanning the null space of A
        a_59 = torch.ones(x1.shape, device=dev)
        A_s = torch.stack(
            (torch.mul(x1, x2), torch.mul(x1, y2), x1,
             torch.mul(y1, x2), torch.mul(y1, y2), y1,
             x2, y2, a_59), dim=-1)

        _, _, v = torch.linalg.svd(A_s)#, full_matrices=False) #A_s # eigenvalues in increasing order

        null_space = v[:, -4:, :].transpose(-1, -2).float()  # the last four rows

        # use the 4 eigenvectors according to the 4 smallest singular values,
        # E is calculated from 4 basis, E = cx*X + cy*Y + cz*Z + cw*W, up to common scale = 1
        null_space_mat = null_space.reshape(null_space.shape[0], 3, 3, null_space.shape[-1]).transpose(1,
                                                                                                       2)  # X, Y, Z, W

        # Step2: expansion of the constraints:
        # determinant constraint det(E) = 0,
        # trace constraint 2EE^TE - trace(EE^T)E = 0
        constraint_mat = self.get_constraint_mat(null_space_mat, batch_size)

        # Step 3: Eliminate part of the matrix to isolate polynomials in z.
        # solve AX=b
        b = constraint_mat[:, :, 10:]
        eliminated_mat = torch.linalg.solve(constraint_mat[:, :, :10], b)

        action_mat = torch.zeros((batch_size, 10, 10), device=dev)
        action_mat[:, 0:3] = eliminated_mat[:, 0:3]
        action_mat[:, 3] = eliminated_mat[:, 4]
        action_mat[:, 4] = eliminated_mat[:, 5]
        action_mat[:, 5] = eliminated_mat[:, 7]
        action_mat[:, 6, 0] = -torch.ones_like(action_mat[:, 6, 0])
        action_mat[:, 7, 1] = -torch.ones_like(action_mat[:, 7, 1])
        action_mat[:, 8, 3] = -torch.ones_like(action_mat[:, 8, 3])
        action_mat[:, 9, 6] = -torch.ones_like(action_mat[:, 9, 6])
        ee, vv = torch.linalg.eig(action_mat)#, eigenvectors=True)# svd

        # put the cx, cy, cz back to get a valid essential matrix
        E_models = null_space.matmul(vv.real[:, -4:])#torch.stack(#.real
        real_mask = (vv.imag[:, -4:].abs() < 1e-3).all(dim=1).flatten()
        E_models = E_models.transpose(-1, -2).reshape(-1, 3, 3).transpose(-1, -2)

        return E_models[real_mask]
//  [python]
def svd_merge(up_1, down_1, up_2, down_2, rank, threshold, device=None):
    org_device = up_1.device
    org_dtype = up_1.dtype

    up_1 = up_1.to(device)
    down_1 = down_1.to(device)
    r_1 = up_1.shape[1]
    weight_1 = up_1.view(-1, r_1) @ down_1.view(r_1, -1)

    if up_2 is not None:
        up_2 = up_2.to(device)
        down_2 = down_2.to(device)
        r_2 = up_2.shape[1]
        weight_2 = up_2.view(-1, r_2) @ down_2.view(r_2, -1)
        weight = weight_1 / r_1 + weight_2 / r_2
    else:
        weight = weight_1 / r_1

    weight = weight.to(dtype=torch.float32) # SVD only supports float32

    U, S, Vh = torch.svd(weight)

    if threshold < 1:
        rank = index_sv_fro(S, threshold) + 1

    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, CLAMP_QUANTILE)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if down_1.dim() == 4:
        U = U.reshape(up_1.shape[0], rank, 1, 1)
        Vh = Vh.reshape(rank, down_1.shape[1], down_1.shape[2], down_1.shape[3])

    up = U.to(org_device, dtype=org_dtype) * math.sqrt(rank)
    down = Vh.to(org_device, dtype=org_dtype) * math.sqrt(rank)

    return up, down
//  [python]
def svd_merge(up_1, down_1, up_2, down_2, rank, threshold, device=None):
    org_device = up_1.device
    org_dtype = up_1.dtype

    up_1 = up_1.to(device)
    down_1 = down_1.to(device)
    r_1 = up_1.shape[1]
    weight_1 = up_1.view(-1, r_1) @ down_1.view(r_1, -1)

    if up_2 is not None:
        up_2 = up_2.to(device)
        down_2 = down_2.to(device)
        r_2 = up_2.shape[1]
        weight_2 = up_2.view(-1, r_2) @ down_2.view(r_2, -1)
        weight = weight_1 / r_1 + weight_2 / r_2
    else:
        weight = weight_1 / r_1

    weight = weight.to(dtype=torch.float32) # SVD only supports float32

    U, S, Vh = torch.linalg.svd(weight)

    if threshold < 1:
        rank = index_sv_fro(S, threshold) + 1

    U = U[:, :rank]
    S = S[:rank]
    U = U @ torch.diag(S)

    Vh = Vh[:rank, :]

    dist = torch.cat([U.flatten(), Vh.flatten()])
    hi_val = torch.quantile(dist, CLAMP_QUANTILE)
    low_val = -hi_val

    U = U.clamp(low_val, hi_val)
    Vh = Vh.clamp(low_val, hi_val)

    if down_1.dim() == 4:
        U = U.reshape(up_1.shape[0], rank, 1, 1)
        Vh = Vh.reshape(rank, down_1.shape[1], down_1.shape[2], down_1.shape[3])

    up = U.to(org_device, dtype=org_dtype) * math.sqrt(rank)
    down = Vh.to(org_device, dtype=org_dtype) * math.sqrt(rank)

    return up, down
//  [python]
def extract_conv(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
        is_cp=False,
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape

    U, S, Vh = torch.linalg.svd(weight.reshape(out_ch, -1))

    if mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == 'quantile' or mode == 'percentile':
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2 and not is_cp:
        return weight, 'full'

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'
//  [python]
def extract_conv(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
        is_cp=False,
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape

    U, S, Vh = linalg.svd(weight.reshape(out_ch, -1))

    if mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == 'quantile' or mode == 'percentile':
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2 and not is_cp:
        return weight, 'full'

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'
//  [python]
def proj2manifold(x):
    u, _, vT = torch.svd(x)
    return u @ vT
//  [python]
def proj2manifold(x):
    u, _, vT = torch.linalg.svd(x)
    return u @ vT
//  [python]
def extract_linear(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch = weight.shape

    U, S, Vh = torch.linalg.svd(weight)

    if mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == 'quantile' or mode == 'percentile':
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2:
        return weight, 'full'

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    diff = (weight - U @ Vh).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'
//  [python]
def extract_linear(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch = weight.shape

    U, S, Vh = linalg.svd(weight)

    if mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == 'quantile' or mode == 'percentile':
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2:
        return weight, 'full'

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    diff = (weight - U @ Vh).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'
//  [python]
def _fit(self, embeddings: torch.Tensor) -> None:
        """
        Perform the PCA. Evaluate ``components``, ``expoained_variance``, ``explained_variance_ratio``,
        ``singular_values``, ``mean``.

        Args:
            embeddings: Embeddings matrix with the shape of ``[n_embeddings, embeddings_dim]``.
        """
        # The code below mirrors code from scikit-learn repository:
        # https://github.com/scikit-learn/scikit-learn/blob/f3f51f9b6/sklearn/decomposition/_pca.py#L491
        n_samples = embeddings.shape[0]
        self.mean = embeddings.mean(dim=0).unsqueeze(0)
        embeddings = embeddings - self.mean
        # if there are more embeddings than its dimension, then we will not perform full matrices evaluation
        full_matrices = embeddings.shape[0] < embeddings.shape[1]
        U, S, V = torch.svd(embeddings, some=full_matrices)
        self.singular_values = S
        self.components = V
        self.explained_variance = self.singular_values**2 / (n_samples - 1)
        self.explained_variance_ratio = self.explained_variance / self.explained_variance.sum()

        # Make components deterministic.
        # Note. In sklearn this operation is done based on the U matrix of SVD decomposition, and
        # here V matrix is used. So, the components of this class and sklearn.decomposition.PCA could differ in sign.
        # See the following for details:
        # https://github.com/scikit-learn/scikit-learn/blob/f3f51f9b6/sklearn/decomposition/_pca.py#L520
        n_components = self.components.shape[0]
        max_abs_rows = torch.argmax(torch.abs(self.components), dim=1)
        signs = torch.sign(self.components[torch.arange(n_components), max_abs_rows])
        self.components *= signs.unsqueeze(1)