jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_linalg_svd_no_compute_uv(self, device, dtype):
        """
        Test the compute_uv=False case. Here we have a different return type than
        numpy: numpy returns S, we return (empty, S, empty)
        """
        t = torch.randn((10, 11), device=device, dtype=dtype)
        np_t = t.cpu().numpy()

        def is_empty(x):
            return x.numel() == 0 and x.dtype == t.dtype and x.device == t.device

        for full_matrices in (True, False):
            # check linalg.svd vs numpy
            np_s = np.linalg.svd(np_t, full_matrices, compute_uv=False)
            USV = torch.linalg.svd(t, full_matrices, compute_uv=False)
            assert is_empty(USV.U)
            self.assertEqual(USV.S, np_s)
            assert is_empty(USV.V)
            # check linalg.svd vs linalg.svd(out=...)
            out = (torch.empty_like(USV.U), torch.empty_like(USV.S), torch.empty_like(USV.V))
            USV = torch.linalg.svd(t, full_matrices, compute_uv=False, out=out)
            assert USV.U is out[0]
            assert USV.S is out[1]
            assert USV.V is out[2]
            self.assertEqual(USV.S, np_s)
//  [python]
def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
    out_size, in_size = weight.size()
    
    U, S, Vh = torch.svd(weight.to(device))
    
    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
    lora_rank = param_dict["new_rank"]
    
    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]
    
    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
    param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
    del U, S, Vh, weight
    return param_dict
//  [python]
def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1):
    out_size, in_size = weight.size()
    
    U, S, Vh = torch.linalg.svd(weight.to(device))
    
    param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale)
    lora_rank = param_dict["new_rank"]
    
    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]
    
    param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu()
    param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu()
    del U, S, Vh, weight
    return param_dict
//  [python]
def random_square_matrix_of_rank(l, rank, dtype=torch.double, device='cpu'):
    assert rank <= l
    A = torch.randn(l, l, dtype=dtype, device=device)
    u, s, vh = torch.svd(A)
    for i in range(l):
        if i >= rank:
            s[i] = 0
        elif s[i] == 0:
            s[i] = 1
    return (u * s.to(dtype).unsqueeze(-2)) @ vh
//  [python]
def random_square_matrix_of_rank(l, rank, dtype=torch.double, device='cpu'):
    assert rank <= l
    A = torch.randn(l, l, dtype=dtype, device=device)
    u, s, vh = torch.linalg.svd(A, full_matrices=False)
    for i in range(l):
        if i >= rank:
            s[i] = 0
        elif s[i] == 0:
            s[i] = 1
    return (u * s.to(dtype).unsqueeze(-2)) @ vh
//  [python]
def random_well_conditioned_matrix(*shape, dtype, device, mean=1.0, sigma=0.001):
    """
    Returns a random rectangular matrix (batch of matrices)
    with singular values sampled from a Gaussian with
    mean `mean` and standard deviation `sigma`.
    The smaller the `sigma`, the better conditioned
    the output matrix is.
    """
    primitive_dtype = {
        torch.float: torch.float,
        torch.double: torch.double,
        torch.cfloat: torch.float,
        torch.cdouble: torch.double
    }
    x = torch.rand(shape, dtype=dtype, device=device)
    m = x.size(-2)
    n = x.size(-1)
    u, s, vh = torch.svd(x)
    s = (torch.randn(*(shape[:-2] + (min(m, n),)), dtype=primitive_dtype[dtype], device=device) * sigma + mean) \
        .sort(-1, descending=True).values.to(dtype)
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def random_well_conditioned_matrix(*shape, dtype, device, mean=1.0, sigma=0.001):
    """
    Returns a random rectangular matrix (batch of matrices)
    with singular values sampled from a Gaussian with
    mean `mean` and standard deviation `sigma`.
    The smaller the `sigma`, the better conditioned
    the output matrix is.
    """
    primitive_dtype = {
        torch.float: torch.float,
        torch.double: torch.double,
        torch.cfloat: torch.float,
        torch.cdouble: torch.double
    }
    x = torch.rand(shape, dtype=dtype, device=device)
    m = x.size(-2)
    n = x.size(-1)
    u, _, vh = torch.linalg.svd(x, full_matrices=False)
    s = (torch.randn(*(shape[:-2] + (min(m, n),)), dtype=primitive_dtype[dtype], device=device) * sigma + mean) \
        .sort(-1, descending=True).values.to(dtype)
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def compute_kernel_bias(self, sentence_vec):
        '''bert-whitening的torch实现
        '''
        vecs = torch.cat(sentence_vec, dim=0)
        self.bias = -vecs.mean(dim=0, keepdims=True)

        cov = torch.cov(vecs.T)  # 协方差
        u, s, vh = torch.svd(cov)
        W = torch.matmul(u, torch.diag(s**0.5))
        self.kernel = torch.inverse(W.T)
//  [python]
def compute_kernel_bias(self, sentence_vec):
        '''bert-whitening的torch实现
        '''
        vecs = torch.cat(sentence_vec, dim=0)
        self.bias = -vecs.mean(dim=0, keepdims=True)

        cov = torch.cov(vecs.T)  # 协方差
        u, s, vh = torch.linalg.svd(cov)
        W = torch.matmul(u, torch.diag(s**0.5))
        self.kernel = torch.linalg.inv(W.T)
//  [python]
def fit(self, source: torch.Tensor, target: torch.Tensor) -> None:
        assert source.shape == target.shape

        # find mean row wise
        centroid_source = torch.mean(source, dim=1)
        centroid_target = torch.mean(target, dim=1)

        # Center the data along the centroid target
        source_m = source - centroid_target.unsqueeze(dim=1)
        target_m = target - centroid_target.unsqueeze(dim=1)

        H = torch.bmm(source_m.permute(0, 2, 1), target_m)
        U, S, Vt = torch.svd(H)

        R = torch.bmm(Vt.permute(0, 2, 1), U.permute(0, 2, 1))
        t = -R @ centroid_source.unsqueeze(dim=2) + centroid_target.unsqueeze(dim=2)

        # Store these
        self.R = R
        self.t = t
//  [python]
def fit(self, source: torch.Tensor, target: torch.Tensor) -> None:
        assert source.shape == target.shape

        # find mean row wise
        centroid_source = torch.mean(source, dim=1)
        centroid_target = torch.mean(target, dim=1)

        # Center the data along the centroid target
        source_m = source - centroid_target.unsqueeze(dim=1)
        target_m = target - centroid_target.unsqueeze(dim=1)

        H = torch.bmm(source_m.permute(0, 2, 1), target_m)
        U, S, Vt = torch.linalg.svd(H)

        R = torch.bmm(Vt.permute(0, 2, 1), U.permute(0, 2, 1))
        t = -R @ centroid_source.unsqueeze(dim=2) + centroid_target.unsqueeze(dim=2)

        # Store these
        self.R = R
        self.t = t
//  [python]
def make_fullrank_matrices_with_distinct_singular_values(*shape, device, dtype, requires_grad=False):
    with torch.no_grad():
        t = make_tensor(shape, device=device, dtype=dtype)
        u, _, vh = torch.svd(t)
        real_dtype = t.real.dtype if t.dtype.is_complex else t.dtype
        k = min(shape[-1], shape[-2])
        # We choose the singular values to be "around one"
        # This is to make the matrix well conditioned
        # s = [2, 3, ..., k+1]
        s = torch.arange(2, k + 2, dtype=real_dtype, device=device)
        # s = [2, -3, 4, ..., (-1)^k k+1]
        s[1::2] *= -1.
        # 1 + 1/s so that the singular values are in the range [2/3, 3/2]
        # This gives a condition number of 9/4, which should be good enough
        s.reciprocal_().add_(1.)
        # Note that the singular values need not be ordered in an SVD so
        # we don't need need to sort S
        x = (u * s.to(u.dtype)) @ vh
    x.requires_grad_(requires_grad)
    return x
//  [python]
def make_fullrank_matrices_with_distinct_singular_values(*shape, device, dtype, requires_grad=False):
    with torch.no_grad():
        t = make_tensor(shape, device=device, dtype=dtype)
        u, _, vh = torch.linalg.svd(t, full_matrices=False)
        real_dtype = t.real.dtype if t.dtype.is_complex else t.dtype
        k = min(shape[-1], shape[-2])
        # We choose the singular values to be "around one"
        # This is to make the matrix well conditioned
        # s = [2, 3, ..., k+1]
        s = torch.arange(2, k + 2, dtype=real_dtype, device=device)
        # s = [2, -3, 4, ..., (-1)^k k+1]
        s[1::2] *= -1.
        # 1 + 1/s so that the singular values are in the range [2/3, 3/2]
        # This gives a condition number of 9/4, which should be good enough
        s.reciprocal_().add_(1.)
        # Note that the singular values need not be ordered in an SVD so
        # we don't need need to sort S
        x = (u * s.to(u.dtype)) @ vh
    x.requires_grad_(requires_grad)
    return x
//  [python]
def random_matrix(rows, columns, *batch_dims, **kwargs):
    """Return rectangular matrix or batches of rectangular matrices.

    Parameters:
      dtype - the data type
      device - the device kind
      singular - when True, the output will be singular
    """
    dtype = kwargs.get('dtype', torch.double)
    device = kwargs.get('device', 'cpu')
    silent = kwargs.get("silent", False)
    singular = kwargs.get("singular", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(rows, columns, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device)
    if A.numel() == 0:
        return A
    u, s, vh = torch.svd(A, some=False)
    k = min(rows, columns)
    s = torch.linspace(1 / (k + 1), 1, k, dtype=dtype, device=device)
    if singular:
        # make matrix singular
        s[k - 1] = 0
        if k > 2:
            # increase the order of singularity so that the pivoting
            # in LU factorization will be non-trivial
            s[0] = 0
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def random_matrix(rows, columns, *batch_dims, **kwargs):
    """Return rectangular matrix or batches of rectangular matrices.

    Parameters:
      dtype - the data type
      device - the device kind
      singular - when True, the output will be singular
    """
    dtype = kwargs.get('dtype', torch.double)
    device = kwargs.get('device', 'cpu')
    silent = kwargs.get("silent", False)
    singular = kwargs.get("singular", False)
    if silent and not torch._C.has_lapack:
        return torch.ones(rows, columns, dtype=dtype, device=device)

    A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device)
    if A.numel() == 0:
        return A
    u, _, vh = torch.linalg.svd(A, full_matrices=False)
    k = min(rows, columns)
    s = torch.linspace(1 / (k + 1), 1, k, dtype=dtype, device=device)
    if singular:
        # make matrix singular
        s[k - 1] = 0
        if k > 2:
            # increase the order of singularity so that the pivoting
            # in LU factorization will be non-trivial
            s[0] = 0
    return (u * s.unsqueeze(-2)) @ vh
//  [python]
def proj_pca(xs_f: torch.Tensor, xs_b: torch.Tensor, reverse: bool) -> Tuple[torch.Tensor, torch.Tensor]:
    # xs: (batch, T, nx)
    # Only use final timestep of xs_f for PCA.
    batch, T, nx = xs_f.shape

    # (batch * T, nx)
    flat_xsf = xs_f.reshape(-1, *xs_f.shape[2:])
    flat_xsb = xs_b.reshape(-1, *xs_b.shape[2:])

    # Center by subtract mean.
    # (batch, nx)
    if reverse:
        # If reverse, use xs_b[0] instead of xs_f[T]
        final_xs_f = xs_b[:, 0, :]
    else:
        final_xs_f = xs_f[:, -1, :]

    mean_pca_xs = torch.mean(final_xs_f, dim=0, keepdim=True)
    final_xs_f -= mean_pca_xs

    # if batch is too large, it will run out of memory.
    if batch > 200:
        rand_idxs = torch.randperm(batch)[:200]
        final_xs_f = final_xs_f[rand_idxs]

    # U: (batch, k)
    # S: (k, k)
    # VT: (k, nx)
    U, S, VT = torch.svd(final_xs_f)

    # log.info("Singular values of xs_f at final timestep:")
    # log.info(S)

    # Keep the first and last directions.
    VT = VT[:2, :]
    # VT = VT[[0, -1], :]

    assert VT.shape == (2, nx)
    V = VT.T

    # Project both xs_f and xs_b onto V.
    flat_xsf -= mean_pca_xs
    flat_xsb -= mean_pca_xs

    proj_xs_f = flat_xsf @ V
    proj_xs_f = proj_xs_f.reshape(batch, T, *proj_xs_f.shape[1:])

    proj_xs_b = flat_xsb @ V
    proj_xs_b = proj_xs_b.reshape(batch, T, *proj_xs_b.shape[1:])

    return proj_xs_f, proj_xs_b
//  [python]
def proj_pca(xs_f: torch.Tensor, xs_b: torch.Tensor, reverse: bool) -> Tuple[torch.Tensor, torch.Tensor]:
    # xs: (batch, T, nx)
    # Only use final timestep of xs_f for PCA.
    batch, T, nx = xs_f.shape

    # (batch * T, nx)
    flat_xsf = xs_f.reshape(-1, *xs_f.shape[2:])
    flat_xsb = xs_b.reshape(-1, *xs_b.shape[2:])

    # Center by subtract mean.
    # (batch, nx)
    if reverse:
        # If reverse, use xs_b[0] instead of xs_f[T]
        final_xs_f = xs_b[:, 0, :]
    else:
        final_xs_f = xs_f[:, -1, :]

    mean_pca_xs = torch.mean(final_xs_f, dim=0, keepdim=True)
    final_xs_f -= mean_pca_xs

    # if batch is too large, it will run out of memory.
    if batch > 200:
        rand_idxs = torch.randperm(batch)[:200]
        final_xs_f = final_xs_f[rand_idxs]

    # U: (batch, k)
    # S: (k, k)
    # VT: (k, nx)
    U, S, VT = torch.linalg.svd(final_xs_f)

    # log.info("Singular values of xs_f at final timestep:")
    # log.info(S)

    # Keep the first and last directions.
    VT = VT[:2, :]
    # VT = VT[[0, -1], :]

    assert VT.shape == (2, nx)
    V = VT.T

    # Project both xs_f and xs_b onto V.
    flat_xsf -= mean_pca_xs
    flat_xsb -= mean_pca_xs

    proj_xs_f = flat_xsf @ V
    proj_xs_f = proj_xs_f.reshape(batch, T, *proj_xs_f.shape[1:])

    proj_xs_b = flat_xsb @ V
    proj_xs_b = proj_xs_b.reshape(batch, T, *proj_xs_b.shape[1:])

    return proj_xs_f, proj_xs_b
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    # This already takes residue identity into account.

    assert A.shape[1] == B.shape[1]
    num_rows, num_cols = A.shape
    if num_rows != 3:
        raise Exception(f"matrix A is not 3xN, it is {num_rows}x{num_cols}")
    num_rows, num_cols = B.shape
    if num_rows != 3:
        raise Exception(f"matrix B is not 3xN, it is {num_rows}x{num_cols}")


    # find mean column wise: 3 x 1
    centroid_A = torch.mean(A, axis=1, keepdims=True)
    centroid_B = torch.mean(B, axis=1, keepdims=True)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    H = Am @ Bm.T

    # find rotation
    U, S, Vt = torch.svd(H)

    R = Vt.T @ U.T
    # special reflection case
    if torch.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=A.device))
        R = (Vt.T @ SS) @ U.T
    assert math.fabs(torch.det(R) - 1) < 3e-3  # note I had to change this error bound to be higher

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    # This already takes residue identity into account.

    assert A.shape[1] == B.shape[1]
    num_rows, num_cols = A.shape
    if num_rows != 3:
        raise Exception(f"matrix A is not 3xN, it is {num_rows}x{num_cols}")
    num_rows, num_cols = B.shape
    if num_rows != 3:
        raise Exception(f"matrix B is not 3xN, it is {num_rows}x{num_cols}")


    # find mean column wise: 3 x 1
    centroid_A = torch.mean(A, axis=1, keepdims=True)
    centroid_B = torch.mean(B, axis=1, keepdims=True)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    H = Am @ Bm.T

    # find rotation
    U, S, Vt = torch.linalg.svd(H)

    R = Vt.T @ U.T
    # special reflection case
    if torch.linalg.det(R) < 0:
        # print("det(R) < R, reflection detected!, correcting for it ...")
        SS = torch.diag(torch.tensor([1.,1.,-1.], device=A.device))
        R = (Vt.T @ SS) @ U.T
    assert math.fabs(torch.linalg.det(R) - 1) < 3e-3  # note I had to change this error bound to be higher

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def batched_svd_torch(F):
    F_shape = F.shape
    F = F.reshape(-1, 3, 3)

    U, sigma, Vt = torch.svd(F)
    flag = torch.det(U) < 0
    U[flag, :, 2] *= -1
    sigma[flag, 2] *= -1

    flag = torch.det(Vt) < 0
    Vt[flag, 2, :] *= -1
    sigma[flag, 2] *= -1

    U = U.reshape(F_shape)
    sigma = sigma.reshape(F_shape[:-1])
    Vt = Vt.reshape(F_shape)

    return U, sigma, Vt
//  [python]
def batched_svd_torch(F):
    F_shape = F.shape
    F = F.reshape(-1, 3, 3)

    U, sigma, Vt = torch.linalg.svd(F)
    flag = torch.det(U) < 0
    U[flag, :, 2] *= -1
    sigma[flag, 2] *= -1

    flag = torch.det(Vt) < 0
    Vt[flag, 2, :] *= -1
    sigma[flag, 2] *= -1

    U = U.reshape(F_shape)
    sigma = sigma.reshape(F_shape[:-1])
    Vt = Vt.reshape(F_shape)

    return U, sigma, Vt
//  [python]
def compute_svd_embeddings(X, rank=None):
	"""
	- X (N x D)
	- rank (int): rank of the approximation (i.e., size of the embedding)
	"""
	assert type(X)==torch.Tensor
	assert X.shape[0] < X.shape[1]

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

	V = Vh.T

	if rank:
		S = S[:rank]
		V = V[:rank]

	return V, S
//  [python]
def compute_svd_embeddings(X, rank=None):
	"""
	- X (N x D)
	- rank (int): rank of the approximation (i.e., size of the embedding)
	"""
	assert type(X)==torch.Tensor
	assert X.shape[0] < X.shape[1]

	U, S, Vh = torch.linalg.svd(X, full_matrices=False)

	V = Vh.T

	if rank:
		S = S[:rank]
		V = V[:rank]

	return V, S
//  [python]
def chordal_l2_mean(rot_samples: torch.Tensor) -> torch.Tensor:
    """Compute the a single rotation average of many samples.

    Chordal L2 mean for rotations as done in "Rotation averaging" Hartley et al (2013).

    Args:
        rot_samples: rotation matrix samples, shape (N, 3, 3)

    Returns:
        average rotation, shape (3, 3)
    """
    C = torch.sum(rot_samples, dim=0)
    U, S, Vt = torch.svd(C)

    S = U @ Vt
    if torch.det(S) < 0:
        S = (
            U
            @ torch.diag(
                torch.tensor(
                    [1.0, 1.0, -1.0], device=rot_samples.device, dtype=rot_samples.dtype
                )
            )
            @ Vt
        )

    return S
//  [python]
def chordal_l2_mean(rot_samples: torch.Tensor) -> torch.Tensor:
    """Compute the a single rotation average of many samples.

    Chordal L2 mean for rotations as done in "Rotation averaging" Hartley et al (2013).

    Args:
        rot_samples: rotation matrix samples, shape (N, 3, 3)

    Returns:
        average rotation, shape (3, 3)
    """
    C = torch.sum(rot_samples, dim=0)
    U, _, Vt = torch.linalg.svd(C)

    S = U @ Vt
    if torch.linalg.det(S) < 0:
        S = (
            U
            @ torch.diag(
                torch.tensor(
                    [1.0, 1.0, -1.0], device=rot_samples.device, dtype=rot_samples.dtype
                )
            )
            @ Vt
        )

    return S
//  [python]
def calculate_erank(x):
    """
    x size: [batch_size, max_length, dim]
    return: average erank in a batch
    """
    U, S, V = torch.svd(x)
    N = torch.linalg.norm(S, ord=1, dim=1, keepdim=True)
    S = S / N
    erank = torch.exp(torch.sum(-S * torch.log(S), dim=1))
    avg_erank = torch.mean(erank)
    return avg_erank.item()
//  [python]
def calculate_erank(x):
    """
    x size: [batch_size, max_length, dim]
    return: average erank in a batch
    """
    _, S, _ = torch.linalg.svd(x)
    N = torch.linalg.norm(S, ord=1, dim=1, keepdim=True)
    S = S / N
    erank = torch.exp(torch.sum(-S * torch.log(S), dim=1))
    avg_erank = torch.mean(erank)
    return avg_erank.item()
//  [python]
def test_linalg_svd_compute_uv(self, device, dtype):
        """
        Test the default case. Here we have the very same behavior as
        NumPy with compute_uv=True.
        """
        t = torch.randn((10, 11), device=device, dtype=dtype)
        np_t = t.cpu().numpy()
        for full_matrices in (True, False):
            # check linalg.svd vs numpy
            expected = np.linalg.svd(np_t, full_matrices, compute_uv=True)
            actual = t.svd(full_matrices)
            # sign/phase of the singular vectors is not unique and therefore absolute values are compared
            self.assertEqual(abs(actual[0]), abs(expected[0]))
            self.assertEqual(actual[1], expected[1])
            self.assertEqual(abs(actual[2]), abs(expected[2]))
            # check linalg.svd vs linalg.svd(out=...)
            out = (torch.empty_like(actual[0]),
                   torch.empty_like(actual[1]),
                   torch.empty_like(actual[2]))
            out2 = t.svd(full_matrices, out=out)
            self.assertEqual(actual, out)
            self.assertEqual(actual, out2)
//  [python]
def test_linalg_svd_compute_uv(self, device, dtype):
        """
        Test the default case. Here we have the very same behavior as
        NumPy with compute_uv=True.
        """
        t = torch.randn((10, 11), device=device, dtype=dtype)
        np_t = t.cpu().numpy()
        for full_matrices in (True, False):
            # check linalg.svd vs numpy
            expected = np.linalg.svd(np_t, full_matrices, compute_uv=True)
            actual = torch.linalg.svd(t, full_matrices)
            # sign/phase of the singular vectors is not unique and therefore absolute values are compared
            self.assertEqual(abs(actual[0]), abs(expected[0]))
            self.assertEqual(actual[1], expected[1])
            self.assertEqual(abs(actual[2]), abs(expected[2]))
            # check linalg.svd vs linalg.svd(out=...)
            out = (torch.empty_like(actual[0]),
                   torch.empty_like(actual[1]),
                   torch.empty_like(actual[2]))
            out2 = torch.linalg.svd(t, full_matrices, out=out)
            self.assertEqual(actual, out)
            self.assertEqual(actual, out2)
//  [python]
def rigid_transform_P2Q(P, Q):
    '''
    mapping from P to Q
    P: [batch, N, 3]
    Q: [batch, N, 3]
    Q = (R @ P.permute(0,2,1)).permute(0,2,1) + t
    '''
    # find mean column wise: 3 x 1
    centroid_P = torch.mean(P, dim=-2, keepdims=True)
    centroid_Q = torch.mean(Q, dim=-2, keepdims=True)

    # subtract mean
    Pm = P - centroid_P
    Qm = Q - centroid_Q

    H = Pm.permute(0,2,1) @ Qm

    # find rotation
    U, S, Vt = torch.svd(H)

    d = torch.sign(torch.det(Vt.permute(0,2,1) @ U.permute(0,2,1)))
    SS = torch.diag_embed(torch.stack([torch.ones_like(d), torch.ones_like(d), d], dim=1))

    R = (Vt.permute(0,2,1) @ SS) @ U.permute(0,2,1)

    t = -(R @ centroid_P.permute(0,2,1)).permute(0,2,1) + centroid_Q

    return R, t
//  [python]
def rigid_transform_P2Q(P, Q):
    '''
    mapping from P to Q
    P: [batch, N, 3]
    Q: [batch, N, 3]
    Q = (R @ P.permute(0,2,1)).permute(0,2,1) + t
    '''
    # find mean column wise: 3 x 1
    centroid_P = torch.mean(P, dim=-2, keepdims=True)
    centroid_Q = torch.mean(Q, dim=-2, keepdims=True)

    # subtract mean
    Pm = P - centroid_P
    Qm = Q - centroid_Q

    H = Pm.permute(0,2,1) @ Qm

    # find rotation
    U, S, Vt = torch.linalg.svd(H)

    d = torch.sign(torch.linalg.det(Vt.permute(0,2,1) @ U.permute(0,2,1)))
    SS = torch.diag_embed(torch.stack([torch.ones_like(d), torch.ones_like(d), d], dim=1))

    R = (Vt.permute(0,2,1) @ SS) @ U.permute(0,2,1)

    t = -(R @ centroid_P.permute(0,2,1)).permute(0,2,1) + centroid_Q

    return R, t
//  [python]
def _compute_delta_dof_pos(self, delta_pose: torch.Tensor, jacobian: torch.Tensor) -> torch.Tensor:
        """Computes the change in dos-position that yields the desired change in pose.

        The method uses the Jacobian mapping from joint-space velocities to end-effector velocities
        to compute the delta-change in the joint-space that moves the robot closer to a desired end-effector
        position.

        Args:
            delta_pose (torch.Tensor): The desired delta pose in shape [N, 3 or 6].
            jacobian (torch.Tensor): The geometric jacobian matrix in shape [N, 3 or 6, num-dof]

        Returns:
            torch.Tensor: The desired delta in joint space.
        """
        if self.cfg.ik_method == "pinv":  # Jacobian pseudo-inverse
            # parameters
            k_val = self._ik_params["k_val"]

            # computation
            jacobian_pinv = torch.linalg.pinv(jacobian)
            delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        elif self.cfg.ik_method == "svd":  # adaptive SVD
            # parameters
            k_val = self._ik_params["k_val"]
            min_singular_value = self._ik_params["min_singular_value"]

            # computation
            # U: 6xd, S: dxd, V: d x num-dof
            U, S, Vh = torch.svd(jacobian)
            S_inv = 1.0 / S
            S_inv = torch.where(S > min_singular_value, S_inv, torch.zeros_like(S_inv))
            jacobian_pinv = (
                torch.transpose(Vh, dim0=1, dim1=2)[:, :, :6]
                @ torch.diag_embed(S_inv)
                @ torch.transpose(U, dim0=1, dim1=2)
            )
            delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        elif self.cfg.ik_method == "trans":  # Jacobian transpose
            # parameters
            k_val = self._ik_params["k_val"]

            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            delta_dof_pos = k_val * jacobian_T @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        elif self.cfg.ik_method == "dls":  # damped least squares
            # parameters
            lambda_val = self._ik_params["lambda_val"]

            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            lambda_matrix = (lambda_val**2) * torch.eye(n=jacobian.shape[1], device=self._device)
            delta_dof_pos = jacobian_T @ torch.inverse(jacobian @ jacobian_T + lambda_matrix) @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        else:
            raise ValueError(f"Unsupported inverse-kinematics method: {self.cfg.ik_method}")

        return delta_dof_pos
//  [python]
def _compute_delta_dof_pos(self, delta_pose: torch.Tensor, jacobian: torch.Tensor) -> torch.Tensor:
        """Computes the change in dos-position that yields the desired change in pose.

        The method uses the Jacobian mapping from joint-space velocities to end-effector velocities
        to compute the delta-change in the joint-space that moves the robot closer to a desired end-effector
        position.

        Args:
            delta_pose (torch.Tensor): The desired delta pose in shape [N, 3 or 6].
            jacobian (torch.Tensor): The geometric jacobian matrix in shape [N, 3 or 6, num-dof]

        Returns:
            torch.Tensor: The desired delta in joint space.
        """
        if self.cfg.ik_method == "pinv":  # Jacobian pseudo-inverse
            # parameters
            k_val = self._ik_params["k_val"]

            # computation
            jacobian_pinv = torch.linalg.pinv(jacobian)
            delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        elif self.cfg.ik_method == "svd":  # adaptive SVD
            # parameters
            k_val = self._ik_params["k_val"]
            min_singular_value = self._ik_params["min_singular_value"]

            # computation
            # U: 6xd, S: dxd, V: d x num-dof
            U, S, Vh = torch.linalg.svd(jacobian)
            S_inv = 1.0 / S
            S_inv = torch.where(S > min_singular_value, S_inv, torch.zeros_like(S_inv))
            jacobian_pinv = (
                torch.transpose(Vh, dim0=1, dim1=2)[:, :, :6]
                @ torch.diag_embed(S_inv)
                @ torch.transpose(U, dim0=1, dim1=2)
            )
            delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        elif self.cfg.ik_method == "trans":  # Jacobian transpose
            # parameters
            k_val = self._ik_params["k_val"]

            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            delta_dof_pos = k_val * jacobian_T @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        elif self.cfg.ik_method == "dls":  # damped least squares
            # parameters
            lambda_val = self._ik_params["lambda_val"]

            # computation
            jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
            lambda_matrix = (lambda_val**2) * torch.eye(n=jacobian.shape[1], device=self._device)
            delta_dof_pos = jacobian_T @ torch.inverse(jacobian @ jacobian_T + lambda_matrix) @ delta_pose.unsqueeze(-1)
            delta_dof_pos = delta_dof_pos.squeeze(-1)

        else:
            raise ValueError(f"Unsupported inverse-kinematics method: {self.cfg.ik_method}")

        return delta_dof_pos
//  [python]
def cal_correlation_svd(bcov):
    _, S, _ = torch.svd(bcov)
    S = torch.mean(S, dim=0)
    return S.detach().cpu().numpy()
//  [python]
def cal_correlation_svd(bcov):
    _, S, _ = torch.linalg.svd(bcov)
    S = torch.mean(S, dim=0)
    return S.detach().cpu().numpy()
//  [python]
def compress(delta_ckpt, ckpt, diffuser=False, compression_ratio=0.6, device='cuda'):
    st = torch.load(f'{delta_ckpt}')

    if not diffuser:
        compressed_key = 'state_dict'
        compressed_st = {compressed_key: {}}
        pretrained_st = torch.load(ckpt)['state_dict']
        if 'embed' in st['state_dict']:
            compressed_st['state_dict']['embed'] = st['state_dict']['embed']
            del st['state_dict']['embed']

        st = st['state_dict']
    else:
        from diffusers import StableDiffusionPipeline
        compressed_key = 'unet'
        compressed_st = {compressed_key: {}}
        pretrained_st = StableDiffusionPipeline.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda")
        pretrained_st = pretrained_st.unet.state_dict()
        if 'modifier_token' in st:
            compressed_st['modifier_token'] = st['modifier_token']
        st = st['unet']

    print("getting compression")
    layers = list(st.keys())
    for name in layers:
        if 'to_k' in name or 'to_v' in name:
            W = st[name].to(device)
            Wpretrain = pretrained_st[name].clone().to(device)
            deltaW = W-Wpretrain

            u, s, vt = torch.svd(deltaW.clone())

            explain = 0
            all_ = (s).sum()
            for i, t in enumerate(s):
                explain += t/(all_)
                if explain > compression_ratio:
                    break

            compressed_st[compressed_key][f'{name}'] = {}
            compressed_st[compressed_key][f'{name}']['u'] = (u[:, :i]@torch.diag(s)[:i, :i]).clone()
            compressed_st[compressed_key][f'{name}']['v'] = vt[:i].clone()
        else:
            compressed_st[compressed_key][f'{name}'] = st[name]

    name = delta_ckpt.replace('delta', 'compressed_delta')
    torch.save(compressed_st, f'{name}')
//  [python]
def compress(delta_ckpt, ckpt, diffuser=False, compression_ratio=0.6, device='cuda'):
    st = torch.load(f'{delta_ckpt}')

    if not diffuser:
        compressed_key = 'state_dict'
        compressed_st = {compressed_key: {}}
        pretrained_st = torch.load(ckpt)['state_dict']
        if 'embed' in st['state_dict']:
            compressed_st['state_dict']['embed'] = st['state_dict']['embed']
            del st['state_dict']['embed']

        st = st['state_dict']
    else:
        from diffusers import StableDiffusionPipeline
        compressed_key = 'unet'
        compressed_st = {compressed_key: {}}
        pretrained_st = StableDiffusionPipeline.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda")
        pretrained_st = pretrained_st.unet.state_dict()
        if 'modifier_token' in st:
            compressed_st['modifier_token'] = st['modifier_token']
        st = st['unet']

    print("getting compression")
    layers = list(st.keys())
    for name in layers:
        if 'to_k' in name or 'to_v' in name:
            W = st[name].to(device)
            Wpretrain = pretrained_st[name].clone().to(device)
            deltaW = W-Wpretrain

            u, s, vt = torch.linalg.svd(deltaW.clone())

            explain = 0
            all_ = (s).sum()
            for i, t in enumerate(s):
                explain += t/(all_)
                if explain > compression_ratio:
                    break

            compressed_st[compressed_key][f'{name}'] = {}
            compressed_st[compressed_key][f'{name}']['u'] = (u[:, :i]@torch.diag(s)[:i, :i]).clone()
            compressed_st[compressed_key][f'{name}']['v'] = vt[:i].clone()
        else:
            compressed_st[compressed_key][f'{name}'] = st[name]

    name = delta_ckpt.replace('delta', 'compressed_delta')
    torch.save(compressed_st, f'{name}')
//  [python]
def whitening_from_covariance(CC):
    """ whitening matrix for a covariance matrix CC
        this is the so-called ZCA whitening matrix
    """
    E,D,V =  torch.svd(CC)
    eps = 1e-6
    Wrot =(E / (D+eps)**.5) @ E.T
    return Wrot
//  [python]
def whitening_from_covariance(CC):
    """ whitening matrix for a covariance matrix CC
        this is the so-called ZCA whitening matrix
    """
    E,D,V =  torch.linalg.svd(CC)
    eps = 1e-6
    Wrot =(E / (D+eps)**.5) @ E.T
    return Wrot
//  [python]
def update_projections(self):
        with torch.no_grad():
            for name, param in self.model.named_parameters():
                if param.requires_grad:
                    grad = param.grad
                    U, S, Vt = torch.svd(grad)
                    self.P[name] = U[:, : self.rank]
                    self.Q[name] = Vt[: self.rank, :].t()
//  [python]
def update_projections(self):
        with torch.no_grad():
            for name, param in self.model.named_parameters():
                if param.requires_grad:
                    grad = param.grad
                    U, _, Vt = torch.linalg.svd(grad, full_matrices=False)
                    self.P[name] = U[:, : self.rank]
                    self.Q[name] = Vt[: self.rank, :].t()
//  [python]
def get_lowrank_tuple_torch_cpu(tensor, max_rank):
	t       = tensor.float()
	u, s, v = torch.svd(t)
	u, s, v = u[:,:max_rank], s[:max_rank], v[:max_rank, :]
	l       = torch.matmul(u, torch.diag(s))
	del t, u, s
	A, B = (v.t(), l.t()) #tensor.t() ~ AB
	return A, B
//  [python]
def get_lowrank_tuple_torch_cpu(tensor, max_rank):
	t       = tensor.float()
	u, s, v = torch.linalg.svd(t)
	u, s, v = u[:,:max_rank], s[:max_rank], v[:max_rank, :]
	l       = torch.matmul(u, torch.diag(s))
	del t, u, s
	A, B = (v.t(), l.t()) #tensor.t() ~ AB
	return A, B
//  [python]
def get_lowrank_tuple_torch_gpu(tensor, max_rank):
	t       = tensor.float().to('cuda')
	u, s, v = torch.svd(t)
	u, s, v = u.to('cpu'), s.to('cpu'), v.to('cpu')
	u, s, v = u[:,:max_rank], s[:max_rank], v[:max_rank, :]
	l       = torch.matmul(u, torch.diag(s))
	del t, u, s
	A, B = (v.t(), l.t()) #tensor.t() ~ AB
	return A, B
//  [python]
def get_lowrank_tuple_torch_gpu(tensor, max_rank):
	t       = tensor.float().to('cuda')
	u, s, v = torch.linalg.svd(t)
	u, s, v = u.to('cpu'), s.to('cpu'), v.to('cpu')
	u, s, v = u[:,:max_rank], s[:max_rank], v[:max_rank, :]
	l       = torch.matmul(u, torch.diag(s))
	del t, u, s
	A, B = (v.t(), l.t()) #tensor.t() ~ AB
	return A, B
//  [python]
def get_euclidean_kabsch(pos, ref, pos_mask):
    #pos [N,M,3]
    #ref [N,M,3]
    #pos_mask [N,M]
    #N : number of examples
    #M : number of atoms
    # R,T maps local reference onto global pos
    if pos_mask is None:
        pos_mask = torch.ones(pos.shape[:2], device=pos.device)
    else:
        if pos_mask.shape[0] != pos.shape[0]:
            raise ValueError("pos_mask should have same number of rows as number of input vectors.")
        if pos_mask.shape[1] != pos.shape[1]:
            raise ValueError("pos_mask should have same number of cols as number of input vector dimensions.")
        if pos_mask.ndim != 2:
            raise ValueError("pos_mask should be 2 dimensional.")

    #Center point clouds
    denom = torch.sum(pos_mask, dim=1, keepdim=True)
    denom[denom==0] = 1.
    pos_mu = torch.sum(pos * pos_mask[:,:,None], dim=1, keepdim=True) / denom[:,:,None]
    ref_mu = torch.sum(ref * pos_mask[:,:,None], dim=1, keepdim=True) / denom[:,:,None]
    pos_c = pos - pos_mu
    ref_c = ref - ref_mu

    #Covariance matrix
    H = torch.einsum('bji,bjk->bik', ref_c, pos_mask[:,:,None] * pos_c)

    U, S, Vh = torch.svd(H)
    #Decide whether we need to correct rotation matrix to ensure right-handed coord system
    locs = torch.det(U @ Vh) < 0
    S[locs,-1] = -S[locs,-1]
    U[locs,:,-1] = -U[locs,:,-1]
    #Rotation matrix    
    R = torch.einsum('bji,bkj->bik',Vh,U) 

    #Translation vector
    T = pos_mu - torch.einsum('bij,bkj->bki',R,ref_mu)
    return T.squeeze(1), R
//  [python]
def get_euclidean_kabsch(pos, ref, pos_mask):
    #pos [N,M,3]
    #ref [N,M,3]
    #pos_mask [N,M]
    #N : number of examples
    #M : number of atoms
    # R,T maps local reference onto global pos
    if pos_mask is None:
        pos_mask = torch.ones(pos.shape[:2], device=pos.device)
    else:
        if pos_mask.shape[0] != pos.shape[0]:
            raise ValueError("pos_mask should have same number of rows as number of input vectors.")
        if pos_mask.shape[1] != pos.shape[1]:
            raise ValueError("pos_mask should have same number of cols as number of input vector dimensions.")
        if pos_mask.ndim != 2:
            raise ValueError("pos_mask should be 2 dimensional.")

    #Center point clouds
    denom = torch.sum(pos_mask, dim=1, keepdim=True)
    denom[denom==0] = 1.
    pos_mu = torch.sum(pos * pos_mask[:,:,None], dim=1, keepdim=True) / denom[:,:,None]
    ref_mu = torch.sum(ref * pos_mask[:,:,None], dim=1, keepdim=True) / denom[:,:,None]
    pos_c = pos - pos_mu
    ref_c = ref - ref_mu

    #Covariance matrix
    H = torch.einsum('bji,bjk->bik', ref_c, pos_mask[:,:,None] * pos_c)

    U, S, Vh = torch.linalg.svd(H)
    #Decide whether we need to correct rotation matrix to ensure right-handed coord system
    locs = torch.linalg.det(U @ Vh) < 0
    S[locs,-1] = -S[locs,-1]
    U[locs,:,-1] = -U[locs,:,-1]
    #Rotation matrix    
    R = torch.einsum('bji,bkj->bik',Vh,U) 

    #Translation vector
    T = pos_mu - torch.einsum('bij,bkj->bki',R,ref_mu)
    return T.squeeze(1), R
//  [python]
def extract_linear(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
) -> Tuple[Tensor, Tensor, int, Tensor]:
    weight = weight.to(device)
    out_ch, in_ch = weight.shape

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

    if mode == 'percentage':
        assert 0 <= mode_param <= 1  # Ensure it's a valid percentage.
        desired_params = mode_param * out_ch * in_ch
        # Solve for lora_rank from the equation
        lora_rank = int(desired_params / (in_ch + out_ch))
    elif mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param).item()
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s).item()
    elif mode == 'quantile':
        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).item()
    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:
        # print(f"rank is higher than it should be")
        lora_rank = int(out_ch / 2)
        # return weight, 'full'
        # print(f"Skipping layer as determined rank is too high")
        # return None, None, None, None

    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, lora_rank, diff
//  [python]
def extract_linear(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
) -> Tuple[Tensor, Tensor, int, Tensor]:
    weight = weight.to(device)
    out_ch, in_ch = weight.shape

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

    if mode == 'percentage':
        assert 0 <= mode_param <= 1  # Ensure it's a valid percentage.
        desired_params = mode_param * out_ch * in_ch
        # Solve for lora_rank from the equation
        lora_rank = int(desired_params / (in_ch + out_ch))
    elif mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param).item()
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s).item()
    elif mode == 'quantile':
        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).item()
    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:
        # print(f"rank is higher than it should be")
        lora_rank = int(out_ch / 2)
        # return weight, 'full'
        # print(f"Skipping layer as determined rank is too high")
        # return None, None, None, None

    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, lora_rank, diff
//  [python]
def extract_conv(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu'
) -> Tuple[Tensor, Tensor, int, Tensor]:
    weight = weight.to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape

    U, S, Vh = torch.svd(weight.reshape(out_ch, -1))
    if mode == 'percentage':
        assert 0 <= mode_param <= 1  # Ensure it's a valid percentage.
        original_params = out_ch * in_ch * kernel_size * kernel_size
        desired_params = mode_param * original_params
        # Solve for lora_rank from the equation
        lora_rank = int(desired_params / (in_ch * kernel_size * kernel_size + out_ch))
    elif mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param).item()
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s).item()
    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).item()
    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:
        lora_rank = int(out_ch / 2)
        print(f"rank is higher than it should be")
        # print(f"Skipping layer as determined rank is too high")
        # return None, None, None, None
        # 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, lora_rank, diff
//  [python]
def extract_conv(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu'
) -> Tuple[Tensor, Tensor, int, Tensor]:
    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 == 'percentage':
        assert 0 <= mode_param <= 1  # Ensure it's a valid percentage.
        original_params = out_ch * in_ch * kernel_size * kernel_size
        desired_params = mode_param * original_params
        # Solve for lora_rank from the equation
        lora_rank = int(desired_params / (in_ch * kernel_size * kernel_size + out_ch))
    elif mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param).item()
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s).item()
    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).item()
    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:
        lora_rank = int(out_ch / 2)
        print(f"rank is higher than it should be")
        # print(f"Skipping layer as determined rank is too high")
        # return None, None, None, None
        # 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, lora_rank, diff