jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def coral(source_features, val_target_features, unlabeled_target_features, labeled_target_features, k, save_path):
    x_s = torch.cat((source_features.to(device), labeled_target_features.to(device)))
    x_t = unlabeled_target_features.to(device)
    x_tv = val_target_features.to(device)
    
    # PCA
    x = torch.cat((x_s, x_t))
    x = x - x.mean(0)
    
    # below line takes forever, save it
    if os.path.isfile(save_path):
        print('Using saved SVD decomposition')
        U, S, Vh = torch.load(save_path)
        U = U.to(device)
        S = S.to(device)
    else:
        U, S, Vh = torch.linalg.svd(x, full_matrices=False)
        torch.save((U, S, Vh), save_path)
    
    x = U[:,:k] @ torch.diag(S[:k])
    x_s = x[:x_s.shape[0]]
    x_t = x[x_s.shape[0]:]
    x_tv = x_tv[:,:k] # NOT GOOD, but I don't need validation for this experiment
    
    x_s_n = x_s - x_s.mean(0) # centered source
    x_t_n = x_t - x_t.mean(0) # centered target
    x_tv = x_tv - x_t.mean(0)

    x_s_cov = torch.matmul(x_s_n.T, x_s_n) / (x_s_n.shape[0] - 1.)
    x_s_cov = x_s_cov + 0.01 * torch.eye(x_s_cov.shape[0]).to(device)
    x_t_cov = torch.matmul(x_t_n.T, x_t_n) / (x_t_n.shape[0] - 1.)
    x_t_cov = x_t_cov + 0.01 * torch.eye(x_t_cov.shape[0]).to(device)

    x_s_cov_sqrt = torch.tensor(sqrtm(x_s_cov.cpu())).to(device)
    x_s_cov_sqrt_inv = x_s_cov_sqrt.inverse()
    x_s_whitened = torch.matmul(x_s_n, x_s_cov_sqrt_inv.float()) # whiten
    x_t_cov_sqrt = torch.tensor(sqrtm(x_t_cov.cpu())).to(device)
    x_s = torch.matmul(x_s_whitened, x_t_cov_sqrt.float()) # recolor with target variance

    x_tu = x_t_n # centered target
    x_t = x_s[source_features.shape[0]:] # target
    x_s = x_s[:source_features.shape[0]] # source
    
    # target unlabeled, target labeled, source, target validation
    return x_tu, x_t, x_s, x_tv
//  [python]
def test_sde_residual(spectral_marginal_sde_instance):
    model = spectral_marginal_sde_instance
    t = torch.linspace(0, 1, 10, requires_grad=True)
    n_reparam = model.diffusion_prior.n_reparam_samples
    z = model.generate_samples(t, n_reparam)
    # this should actually have an additional batch dimension
    f = model.drift(t, z) + torch.randn(n_reparam, 10, model.d) * 1e-2
    r = model.unweighted_residual_loss(model.drift(t, z), f)
    # checking shape
    assert r.shape == (n_reparam, 10)
    # checking residual loss
    d = f - model.drift(t, z)
    p_noise = model.diffusion_prior.process_noise
    if p_noise.dim() == 3:
        r_linalg_solve = (
            d * torch.solve(d.transpose(-2, -1), p_noise).transpose(-2, -1)
        ).sum(-1)
    else:
        U, S, Vh = torch.svd(p_noise)
        solve = (d @ (U * (S / (S.pow(2) + 1e-12)))) @ Vh.T
        r_linalg_solve = torch.einsum("ijk,ijk->ij", solve, d)
    with torch.no_grad():
        npt.assert_array_almost_equal(r, r_linalg_solve, decimal=2)
//  [python]
def test_sde_residual(spectral_marginal_sde_instance):
    model = spectral_marginal_sde_instance
    t = torch.linspace(0, 1, 10, requires_grad=True)
    n_reparam = model.diffusion_prior.n_reparam_samples
    z = model.generate_samples(t, n_reparam)
    # this should actually have an additional batch dimension
    f = model.drift(t, z) + torch.randn(n_reparam, 10, model.d) * 1e-2
    r = model.unweighted_residual_loss(model.drift(t, z), f)
    # checking shape
    assert r.shape == (n_reparam, 10)
    # checking residual loss
    d = f - model.drift(t, z)
    p_noise = model.diffusion_prior.process_noise
    if p_noise.dim() == 3:
        r_linalg_solve = (
            d * torch.linalg.solve(p_noise, d.transpose(-2, -1)).transpose(-2, -1)
        ).sum(-1)
    else:
        U, S, Vh = torch.linalg.svd(p_noise, full_matrices=False)
        solve = (d @ (U * (S / (S.pow(2) + 1e-12)))) @ Vh.T
        r_linalg_solve = torch.einsum("ijk,ijk->ij", solve, d)
    with torch.no_grad():
        npt.assert_array_almost_equal(r, r_linalg_solve, decimal=2)
//  [python]
def perform_svd(self):
        # shape
        weight_reshaped = rearrange(self.weight, 'co cin h w -> co (cin h w)')
        self.U, self.S, self.Vh = torch.svd(weight_reshaped, some=False)
        self.done_svd = True
//  [python]
def perform_svd(self):
        # shape
        weight_reshaped = rearrange(self.weight, 'co cin h w -> co (cin h w)')
        self.U, self.S, self.Vh = torch.linalg.svd(weight_reshaped, full_matrices=False)
        self.done_svd = True
//  [python]
def __init__(
        self, 
        in_channels: int, 
        out_channels: int,
        kernel_size: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
        assert type(kernel_size) is int
        weight_reshaped = rearrange(self.weight, 'co cin h w -> co (cin h w)')
        self.U, self.S, self.Vh = torch.svd(weight_reshaped)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        in_channels: int, 
        out_channels: int,
        kernel_size: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
        assert type(kernel_size) is int
        weight_reshaped = rearrange(self.weight, 'co cin h w -> co (cin h w)')
        self.U, self.S, self.Vh = torch.linalg.svd(weight_reshaped, full_matrices=False)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def batch_kabsch_torch(A, B):
    '''
    A: [B, N, 3]
    B: [B, N, 3]
    '''
    a_mean = A.mean(dim=1, keepdims=True)
    b_mean = B.mean(dim=1, keepdims=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = torch.bmm(A_c.transpose(1,2), B_c)  # [B, 3, 3]
    U, S, Vt = torch.svd(H)  # [B, 3, 3]
    V = Vt.transpose(1, 2)
    # rms
    d = ((torch.det(U) * torch.det(V)) < 0.0).long()  # [B]
    nSS = torch.diag(torch.tensor([1. for _ in range(len(U))], device=U.device, dtype=U.dtype))
    SS = torch.diag(torch.tensor([1. for _ in range(len(U) - 1)] + [-1.], device=U.device, dtype=U.dtype))
    bSS = torch.stack([nSS, SS], dim=0)[d]  # [B, 3, 3]
    U = torch.bmm(U, bSS)
    # Rotation matrix
    R = torch.bmm(V, U.transpose(1,2))  # [B, 3, 3]
    # Translation vector
    t = b_mean - torch.bmm(R, a_mean.transpose(1,2)).transpose(1,2)
    A_aligned = torch.bmm(R, A.transpose(1,2)).transpose(1,2) + t
    return A_aligned, R, t
//  [python]
def batch_kabsch_torch(A, B):
    '''
    A: [B, N, 3]
    B: [B, N, 3]
    '''
    a_mean = A.mean(dim=1, keepdims=True)
    b_mean = B.mean(dim=1, keepdims=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = torch.bmm(A_c.transpose(1,2), B_c)  # [B, 3, 3]
    U, S, Vt = torch.linalg.svd(H)  # [B, 3, 3]
    V = Vt.transpose(1, 2)
    # rms
    d = ((torch.linalg.det(U) * torch.linalg.det(V)) < 0.0).long()  # [B]
    nSS = torch.diag(torch.tensor([1. for _ in range(len(U))], device=U.device, dtype=U.dtype))
    SS = torch.diag(torch.tensor([1. for _ in range(len(U) - 1)] + [-1.], device=U.device, dtype=U.dtype))
    bSS = torch.stack([nSS, SS], dim=0)[d]  # [B, 3, 3]
    U = torch.bmm(U, bSS)
    # Rotation matrix
    R = torch.bmm(V, U.transpose(1,2))  # [B, 3, 3]
    # Translation vector
    t = b_mean - torch.bmm(R, a_mean.transpose(1,2)).transpose(1,2)
    A_aligned = torch.bmm(R, A.transpose(1,2)).transpose(1,2) + t
    return A_aligned, R, t
//  [python]
def get_transformation(hm, region_score, batch, v0, v1):
    with torch.no_grad():
        # hd is the weighted center of ecah part
        hd0 = scatter_sum(hm[:, :, None] * v0[:, None], batch, dim=0)
        hd1 = scatter_sum(hm[:, :, None] * v1[:, None], batch, dim=0)  # (B, 40, 3)

        B, K, _ = hd0.shape                   # 1 40 3
        seg = torch.max(region_score, 1)[1]  # (B*V,)

        R_batch = []
        for i in range(B):
            center0 = hd0[i]  # (40, 3)
            center1 = hd1[i]  # (40, 3)
            seg_single = seg[batch==i]  # (V,)   indicates each vertex belongs to which part
            v0_single = v0[batch==i]  # (V, 3)
            v1_single = v1[batch==i]  # (V, 3)
            seg_max = torch.max(seg_single) + 1

            scatter_center0 = center0[seg_single]
            scatter_center1 = center1[seg_single]

            v0_single = v0_single - scatter_center0
            v1_single = v1_single - scatter_center1

            M = torch.bmm(v0_single[:, :, None], v1_single[:, None])  # (V, 3, 3)
            M_sum = scatter_sum(M, seg_single, dim=0)  # (40, 3, 3)
            U, S, Vh = torch.svd(M_sum)
            R_temp = torch.bmm(U, Vh).permute(0, 2, 1)  # (40, 3, 3)
            R = torch.eye(3)[None].repeat(K, 1, 1).float().to(R_temp.device)
            R[:seg_max] = R_temp
            R = rotation_matrix_to_quaternion(R, order=wxyz)
            R_batch.append(R)

        R_batch = torch.stack(R_batch, 0)  # (B, 40, 4)
        t_batch = hd1 - hd0                # (B, 40, 3)
    return torch.cat((t_batch, R_batch), 2)
//  [python]
def get_transformation(hm, region_score, batch, v0, v1):
    with torch.no_grad():
        # hd is the weighted center of ecah part
        hd0 = scatter_sum(hm[:, :, None] * v0[:, None], batch, dim=0)
        hd1 = scatter_sum(hm[:, :, None] * v1[:, None], batch, dim=0)  # (B, 40, 3)

        B, K, _ = hd0.shape                   # 1 40 3
        seg = torch.max(region_score, 1)[1]  # (B*V,)

        R_batch = []
        for i in range(B):
            center0 = hd0[i]  # (40, 3)
            center1 = hd1[i]  # (40, 3)
            seg_single = seg[batch==i]  # (V,)   indicates each vertex belongs to which part
            v0_single = v0[batch==i]  # (V, 3)
            v1_single = v1[batch==i]  # (V, 3)
            seg_max = torch.max(seg_single) + 1

            scatter_center0 = center0[seg_single]
            scatter_center1 = center1[seg_single]

            v0_single = v0_single - scatter_center0
            v1_single = v1_single - scatter_center1

            M = torch.bmm(v0_single[:, :, None], v1_single[:, None])  # (V, 3, 3)
            M_sum = scatter_sum(M, seg_single, dim=0)  # (40, 3, 3)
            U, S, Vh = torch.linalg.svd(M_sum)
            R_temp = torch.bmm(U, Vh).permute(0, 2, 1)  # (40, 3, 3)
            R = torch.eye(3)[None].repeat(K, 1, 1).float().to(R_temp.device)
            R[:seg_max] = R_temp
            R = rotation_matrix_to_quaternion(R, order=wxyz)
            R_batch.append(R)

        R_batch = torch.stack(R_batch, 0)  # (B, 40, 4)
        t_batch = hd1 - hd0                # (B, 40, 3)
    return torch.cat((t_batch, R_batch), 2)
//  [python]
def compute_svd(self):
        """
        Computes the SVD of the Hankel matrix.
        """

        if self.verbose:
            print("Computing SVD on Hankel matrix ...")
        if self.ntrials > 1: #flatten across trials for 3d
            H = self.H.reshape(self.H.shape[0] * self.H.shape[1], self.H.shape[2])
        else:
            H = self.H
        # compute the SVD
        U, S, Vh = torch.svd(H.T)
        
        # update attributes
        V = Vh.T
        self.U = U
        self.S = S
        self.V = V

        # construct the singuar value matrix and its inverse
        # dim = self.n_delays * self.n
        # s = len(S)
        # self.S_mat = torch.zeros(dim, dim,dtype=torch.float32).to(self.device)
        # self.S_mat_inv = torch.zeros(dim, dim,dtype=torch.float32).to(self.device)
        self.S_mat = torch.diag(S).to(self.device)
        self.S_mat_inv= torch.diag(1 / S).to(self.device)

        # compute explained variance
        exp_variance_inds = self.S**2 / ((self.S**2).sum())
        cumulative_explained = torch.cumsum(exp_variance_inds, 0)
        self.cumulative_explained_variance = cumulative_explained
        
        #make the X and Y components of the regression by staggering the hankel eigen-time delay coordinates by time
        if self.reduced_rank_reg:
            V = self.V
        else:
            V = self.V

        if self.ntrials > 1:
            if V.numel() < self.H.numel():
                raise ValueError("The dimension of the SVD of the Hankel matrix is smaller than the dimension of the Hankel matrix itself. \n \
                                 This is likely due to the number of time points being smaller than the number of dimensions. \n \
                                 Please reduce the number of delays.")
                                 
            V = V.reshape(self.H.shape)

            #first reshape back into Hankel shape, separated by trials
            newshape = (self.H.shape[0]*(self.H.shape[1]-1),self.H.shape[2])
            self.Vt_minus = V[:,:-1].reshape(newshape)
            self.Vt_plus = V[:,1:].reshape(newshape)
        else:
            self.Vt_minus = V[:-1]
            self.Vt_plus = V[1:]


        if self.verbose:
            print("SVD complete!")
//  [python]
def compute_svd(self):
        """
        Computes the SVD of the Hankel matrix.
        """

        if self.verbose:
            print("Computing SVD on Hankel matrix ...")
        if self.ntrials > 1: #flatten across trials for 3d
            H = self.H.reshape(self.H.shape[0] * self.H.shape[1], self.H.shape[2])
        else:
            H = self.H
        # compute the SVD
        U, S, Vh = torch.linalg.svd(H.T, full_matrices=False)
        
        # update attributes
        V = Vh.T
        self.U = U
        self.S = S
        self.V = V

        # construct the singuar value matrix and its inverse
        # dim = self.n_delays * self.n
        # s = len(S)
        # self.S_mat = torch.zeros(dim, dim,dtype=torch.float32).to(self.device)
        # self.S_mat_inv = torch.zeros(dim, dim,dtype=torch.float32).to(self.device)
        self.S_mat = torch.diag(S).to(self.device)
        self.S_mat_inv= torch.diag(1 / S).to(self.device)

        # compute explained variance
        exp_variance_inds = self.S**2 / ((self.S**2).sum())
        cumulative_explained = torch.cumsum(exp_variance_inds, 0)
        self.cumulative_explained_variance = cumulative_explained
        
        #make the X and Y components of the regression by staggering the hankel eigen-time delay coordinates by time
        if self.reduced_rank_reg:
            V = self.V
        else:
            V = self.V

        if self.ntrials > 1:
            if V.numel() < self.H.numel():
                raise ValueError("The dimension of the SVD of the Hankel matrix is smaller than the dimension of the Hankel matrix itself. \n \
                                 This is likely due to the number of time points being smaller than the number of dimensions. \n \
                                 Please reduce the number of delays.")
                                 
            V = V.reshape(self.H.shape)

            #first reshape back into Hankel shape, separated by trials
            newshape = (self.H.shape[0]*(self.H.shape[1]-1),self.H.shape[2])
            self.Vt_minus = V[:,:-1].reshape(newshape)
            self.Vt_plus = V[:,1:].reshape(newshape)
        else:
            self.Vt_minus = V[:-1]
            self.Vt_plus = V[1:]


        if self.verbose:
            print("SVD complete!")
//  [python]
def _svd_lowrank(
    A: Tensor,
    q: Optional[int] = 6,
    niter: Optional[int] = 2,
    M: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor]:
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M = M.broadcast_to(A.size())
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n or n > q:
        # computing the SVD approximation of a transpose in
        # order to keep B shape minimal (the m < n case) or the V
        # shape small (the n > q case)
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        assert B_t.shape[-2] == m, (B_t.shape, m)
        assert B_t.shape[-1] == q, (B_t.shape, q)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.svd(B_t)
        V = Vh.mH
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        B_t = _utils.transpose(B)
        assert B_t.shape[-2] == q, (B_t.shape, q)
        assert B_t.shape[-1] == n, (B_t.shape, n)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.svd(B_t)
        V = Vh.mH
        U = Q.matmul(U)

    return U, S, V
//  [python]
def _svd_lowrank(
    A: Tensor,
    q: Optional[int] = 6,
    niter: Optional[int] = 2,
    M: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor]:
    q = 6 if q is None else q
    m, n = A.shape[-2:]
    matmul = _utils.matmul
    if M is None:
        M_t = None
    else:
        M = M.broadcast_to(A.size())
        M_t = _utils.transpose(M)
    A_t = _utils.transpose(A)

    # Algorithm 5.1 in Halko et al 2009, slightly modified to reduce
    # the number conjugate and transpose operations
    if m < n or n > q:
        # computing the SVD approximation of a transpose in
        # order to keep B shape minimal (the m < n case) or the V
        # shape small (the n > q case)
        Q = get_approximate_basis(A_t, q, niter=niter, M=M_t)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B_t = matmul(A, Q_c)
        else:
            B_t = matmul(A, Q_c) - matmul(M, Q_c)
        assert B_t.shape[-2] == m, (B_t.shape, m)
        assert B_t.shape[-1] == q, (B_t.shape, q)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
        V = Vh.mH
        V = Q.matmul(V)
    else:
        Q = get_approximate_basis(A, q, niter=niter, M=M)
        Q_c = _utils.conjugate(Q)
        if M is None:
            B = matmul(A_t, Q_c)
        else:
            B = matmul(A_t, Q_c) - matmul(M_t, Q_c)
        B_t = _utils.transpose(B)
        assert B_t.shape[-2] == q, (B_t.shape, q)
        assert B_t.shape[-1] == n, (B_t.shape, n)
        assert B_t.shape[-1] <= B_t.shape[-2], B_t.shape
        U, S, Vh = torch.linalg.svd(B_t, full_matrices=False)
        V = Vh.mH
        U = Q.matmul(U)

    return U, S, V
//  [python]
def __init__(
        self, 
        in_features: int, 
        out_features: int, 
        scale: float = 1.0,
        **kwargs
    ):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        self.U, self.S, self.Vh = torch.svd(self.weight)        
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        in_features: int, 
        out_features: int, 
        scale: float = 1.0,
        **kwargs
    ):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        self.U, self.S, self.Vh = torch.linalg.svd(self.weight, full_matrices=False)        
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        in_channels: int, 
        out_channels: int,
        kernel_size: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.Conv1d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
        assert type(kernel_size) is int
        weight_reshaped = rearrange(self.weight, 'co cin h w -> co (cin h w)')
        self.U, self.S, self.Vh = torch.svd(weight_reshaped)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        in_channels: int, 
        out_channels: int,
        kernel_size: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.Conv1d.__init__(self, in_channels, out_channels, kernel_size, **kwargs)
        assert type(kernel_size) is int
        weight_reshaped = rearrange(self.weight, 'co cin h w -> co (cin h w)')
        self.U, self.S, self.Vh = torch.linalg.svd(weight_reshaped, full_matrices=False)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def perform_svd(self):
        self.U, self.S, self.Vh = torch.svd(self.weight, some=False)
        self.done_svd = True
//  [python]
def perform_svd(self):
        self.U, self.S, self.Vh = torch.linalg.svd(self.weight, full_matrices=False)
        self.done_svd = True
//  [python]
def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
        self.U, self.S, self.Vh = torch.svd(self.weight)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self,
        num_embeddings: int,
        embedding_dim: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
        self.U, self.S, self.Vh = torch.linalg.svd(self.weight, full_matrices=False)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        normalized_shape: int, 
        scale: float = 1.0,
        **kwargs
    ):
        nn.LayerNorm.__init__(self, normalized_shape=normalized_shape, **kwargs)
        self.U, self.S, self.Vh = torch.svd(self.weight.unsqueeze(0))
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        normalized_shape: int, 
        scale: float = 1.0,
        **kwargs
    ):
        nn.LayerNorm.__init__(self, normalized_shape=normalized_shape, **kwargs)
        self.U, self.S, self.Vh = torch.linalg.svd(self.weight.unsqueeze(0), full_matrices=False)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        num_groups: int, 
        num_channels: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.GroupNorm.__init__(self, num_groups, num_channels, **kwargs)
        self.U, self.S, self.Vh = torch.svd(self.weight.unsqueeze(0))
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def __init__(
        self, 
        num_groups: int, 
        num_channels: int,
        scale: float = 1.0,
        **kwargs
    ):
        nn.GroupNorm.__init__(self, num_groups, num_channels, **kwargs)
        self.U, self.S, self.Vh = torch.linalg.svd(self.weight.unsqueeze(0), full_matrices=False)
        # initialize to 0 for smooth tuning 
        self.delta = nn.Parameter(torch.zeros_like(self.S))
        self.weight.requires_grad = False
        self.done_svd = False
        self.scale = scale
        self.reset_parameters()
//  [python]
def decompose(self, decomposing_mode: str) -> None:
        """Decomposes the weight matrix in singular value decomposition.
        Replaces the weights with U, S, Vh matrices such that weights = U * S * Vh.
        Args:
            decomposing_mode: ``'channel'`` or ``'spatial'`` weights reshaping method.
        Raises:
            ValueError: If ``decomposing_mode`` not in valid values.
        """
        self.__set_decomposing_params(decomposing_mode=decomposing_mode)
        W = self.weight.permute(self.decomposing['permute']).reshape(
            self.decomposing['decompose_shape'])
        U, S, Vh = torch.svd(W)
        self.U = Parameter(U)
        self.S = Parameter(S)
        self.Vh = Parameter(Vh)
        self.register_parameter('weight', None)
//  [python]
def decompose(self, decomposing_mode: str) -> None:
        """Decomposes the weight matrix in singular value decomposition.
        Replaces the weights with U, S, Vh matrices such that weights = U * S * Vh.
        Args:
            decomposing_mode: ``'channel'`` or ``'spatial'`` weights reshaping method.
        Raises:
            ValueError: If ``decomposing_mode`` not in valid values.
        """
        self.__set_decomposing_params(decomposing_mode=decomposing_mode)
        W = self.weight.permute(self.decomposing['permute']).reshape(
            self.decomposing['decompose_shape'])
        U, S, Vh = torch.linalg.svd(W, full_matrices=False)
        self.U = Parameter(U)
        self.S = Parameter(S)
        self.Vh = Parameter(Vh)
        self.register_parameter('weight', None)
//  [python]
def rotate_kp(kp, neighbors: torch.Tensor, radius):
        with torch.no_grad():
            # m = neighbors.mean(-2, keepdim=True) > 0
            # print("mean", m)
            within = neighbors.norm(dim=-1, keepdim=True) < radius
            n = neighbors * within
            cross_cov = n.transpose(-2, -1) @ n
            U, S, V = torch.svd(cross_cov)
            # V = torch.where(m.transpose(-2,-1).expand_as(V),V,-V)
            V.unsqueeze_(-3)
            return kp @ V
//  [python]
def rotate_kp(kp, neighbors: torch.Tensor, radius):
        with torch.no_grad():
            # m = neighbors.mean(-2, keepdim=True) > 0
            # print("mean", m)
            within = neighbors.norm(dim=-1, keepdim=True) < radius
            n = neighbors * within
            cross_cov = n.transpose(-2, -1) @ n
            _, _, V = torch.linalg.svd(cross_cov)
            # V = torch.where(m.transpose(-2,-1).expand_as(V),V,-V)
            V.unsqueeze_(-3)
            return kp @ V
//  [python]
def cov_kfe_A(cls, module, in_data):
        n, cin_ks, out_size = in_data.shape
        if n * out_size < cin_ks:
            m = in_data.permute(0, 2, 1).view(n * out_size, -1)
            U, S, Vt = torch.svd(m)
            return Vt.T  # (c_in)(kernel_size) x (c_in)(kernel_size)
        else:
            A = cls.cov_kron_A(module, in_data)
            _, U = torch.linalg.eigh(A)
            return U
//  [python]
def cov_kfe_A(cls, module, in_data):
        n, cin_ks, out_size = in_data.shape
        if n * out_size < cin_ks:
            m = in_data.permute(0, 2, 1).view(n * out_size, -1)
            _, _, Vt = torch.linalg.svd(m, full_matrices=True)
            return Vt.T  # (c_in)(kernel_size) x (c_in)(kernel_size)
        else:
            A = cls.cov_kron_A(module, in_data)
            _, U = torch.linalg.eigh(A)
            return U
//  [python]
def cov_kfe_B(cls, module, out_grads):
        n, c_out, out_size = out_grads.shape  # n x c_out x out_size
        if n * out_size < c_out:
            m = out_grads.permute(0, 2, 1).view(n * out_size, -1)
            U, S, V = torch.svd(m)
            return V.T  # c_out x c_out
        else:
            B = cls.cov_kron_B(module, out_grads)
            _, U = torch.linalg.eigh(B)
            return U
//  [python]
def cov_kfe_B(cls, module, out_grads):
        n, c_out, out_size = out_grads.shape  # n x c_out x out_size
        if n * out_size < c_out:
            m = out_grads.permute(0, 2, 1).view(n * out_size, -1)
            _, _, Vt = torch.linalg.svd(m, full_matrices=True)
            return Vt.T  # c_out x c_out
        else:
            B = cls.cov_kron_B(module, out_grads)
            _, U = torch.linalg.eigh(B)
            return U
//  [python]
def random_sketch_and_gram(self, module, in_data, out_grads):
        rank = self._truncated_rank
        if rank is not None and rank < in_data.shape[2]:
            u, s, vt = torch.svd(
                out_grads, some=False)  # n x c_out x out_size
            v = vt.transpose(1, 2)
            u = u[:, :, :rank]  # n x c_out x r
            v = v[:, :, :rank]  # n x out_size x r
            s = s[:, :rank].unsqueeze(1)  # n x 1 x r
            in_data = torch.bmm(in_data, v)  # n x (c_in)(ks) x r
            out_grads = u * s.expand(u.shape)  # n x c_out x r
        return super().random_sketch_and_gram(module, in_data, out_grads)
//  [python]
def random_sketch_and_gram(self, module, in_data, out_grads):
        rank = self._truncated_rank
        if rank is not None and rank < in_data.shape[2]:
            u, s, vt = torch.linalg.svd(
                out_grads, full_matrices=False)  # n x c_out x out_size
            v = vt.transpose(1, 2)
            u = u[:, :, :rank]  # n x c_out x r
            v = v[:, :, :rank]  # n x out_size x r
            s = s[:, :rank].unsqueeze(1)  # n x 1 x r
            in_data = torch.bmm(in_data, v)  # n x (c_in)(ks) x r
            out_grads = u * s.expand(u.shape)  # n x c_out x r
        return super().random_sketch_and_gram(module, in_data, out_grads)
//  [python]
def update_L(self, lam=0.1):
        with torch.no_grad():
            kernel = torch.view_as_complex(self.kernel)
            normalized_kernel = kernel #/ ((kernel * kernel.conj()).sum(dim=(2, 3, 4), keepdim=True) * 2 / self.image_size ** 2).real ** 0.5
            kk = (normalized_kernel.permute((0, 3, 4, 1, 2)).unsqueeze(-1) * normalized_kernel.conj().permute(
                (0, 3, 4, 2, 1)).unsqueeze(3)).sum(dim=-2).permute(0, 3, 4, 1, 2)
            self.L.data = (1-lam) * self.L.data + lam * (torch.svd(kk.float().permute(0, 3, 4, 1, 2))[1].view(self.group, -1).max(dim=-1)[0]).detach()
//  [python]
def update_L(self, lam=0.1):
        with torch.no_grad():
            kernel = torch.view_as_complex(self.kernel)
            normalized_kernel = kernel #/ ((kernel * kernel.conj()).sum(dim=(2, 3, 4), keepdim=True) * 2 / self.image_size ** 2).real ** 0.5
            kk = (normalized_kernel.permute((0, 3, 4, 1, 2)).unsqueeze(-1) * normalized_kernel.conj().permute(
                (0, 3, 4, 2, 1)).unsqueeze(3)).sum(dim=-2).permute(0, 3, 4, 1, 2)
            self.L.data = (1-lam) * self.L.data + lam * (torch.linalg.svd(kk.float().permute(0, 3, 4, 1, 2), full_matrices=False)[1].view(self.group, -1).max(dim=-1)[0]).detach()
//  [python]
def test_unpack2():
    m = variable(torch.eye(10, 10))
    u, s, v = torch.svd(m).unpack(["u", "s", "v"])
    assert torch.equal(u(), torch.eye(10, 10))
    assert torch.equal(s(), torch.ones(10))
    assert torch.equal(v(), torch.eye(10))
//  [python]
def test_unpack2():
    m = variable(torch.eye(10, 10))
    u, s, v = torch.linalg.svd(m).unpack(["u", "s", "v"])
    assert torch.equal(u(), torch.eye(10, 10))
    assert torch.equal(s(), torch.ones(10))
    assert torch.equal(v(), torch.eye(10))
//  [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

    # TODO rank k randomized svd if k small enough
    U, s, Vt = torch.svd(W, some=False)
    Wa = U[:, :k]
    Wb = Vt[:k]

    # scale matrices equally for numerical "load-balancing"
    s_sqrt = torch.sqrt(s[:k])  # s is already a vector, not mat
    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

    # TODO rank k randomized svd if k small enough
    U, s, Vt = torch.linalg.svd(W, full_matrices=False)
    Wa = U[:, :k]
    Wb = Vt[:k]

    # scale matrices equally for numerical "load-balancing"
    s_sqrt = torch.sqrt(s[:k])  # s is already a vector, not mat
    Wa *= s_sqrt
    Wb *= s_sqrt.reshape(-1, 1)
    return Wa, Wb
//  [python]
def SVD_Precoding_tensor(H, P, d):
    U,D,V = torch.svd(H)
    W_svd = V.conj().T[:,:d]
    M_svd = U

    W_svd_norm = torch.sqrt(torch.trace(torch.mm(W_svd,W_svd.conj().T)))   #power norm
    W_svd = W_svd * np.sqrt(P) / W_svd_norm
    return W_svd, D, M_svd
//  [python]
def SVD_Precoding_tensor(H, P, d):
    U,D,V = torch.linalg.svd(H, full_matrices=True)
    W_svd = V.conj().T[:,:d]
    M_svd = U

    W_svd_norm = torch.sqrt(torch.trace(torch.mm(W_svd,W_svd.conj().T)))   #power norm
    W_svd = W_svd * np.sqrt(P) / W_svd_norm
    return W_svd, D, M_svd
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    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.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) < 1e-5

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def rigid_transform_Kabsch_3D_torch(A, B):
    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) < 1e-5

    t = -R @ centroid_A + centroid_B
    return R, t
//  [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 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 Circular_Gaussian(self, Tx_sig, snr=10, Group_enable=True):
        H_raw = 1/math.sqrt(2)* (torch.randn((self.Nr, self.Nt))+1j*torch.randn(self.Nr, self.Nt))    # Nr * Nt
        H_raw = torch.tensor(H_raw).to(self.device)
        sigmaH = self.P * 10**(-snr/10) 
        H_quan = H_raw + math.sqrt(sigmaH/2)* (torch.randn(self.Nr, self.Nt)+1j*torch.randn(self.Nr, self.Nt)).to(self.device)
        # H_quan.to(self.device)
        V,_,_ = SVD_Precoding_tensor(H_quan, self.P, self.d)
        U,D,_ = torch.svd(H_raw)
        if Group_enable:
            return self.Trans_Procedure_group(Tx_sig, H_raw, V, D, U, snr)
        else:
            return self.Trans_Procedure_element(Tx_sig, H_raw, V, D, U, snr)
//  [python]
def Circular_Gaussian(self, Tx_sig, snr=10, Group_enable=True):
        H_raw = 1/math.sqrt(2)* (torch.randn((self.Nr, self.Nt))+1j*torch.randn(self.Nr, self.Nt))    # Nr * Nt
        H_raw = torch.tensor(H_raw).to(self.device)
        sigmaH = self.P * 10**(-snr/10) 
        H_quan = H_raw + math.sqrt(sigmaH/2)* (torch.randn(self.Nr, self.Nt)+1j*torch.randn(self.Nr, self.Nt)).to(self.device)
        # H_quan.to(self.device)
        V,_,_ = SVD_Precoding_tensor(H_quan, self.P, self.d)
        U,D,_ = torch.linalg.svd(H_raw, full_matrices=False)
        if Group_enable:
            return self.Trans_Procedure_group(Tx_sig, H_raw, V, D, U, snr)
        else:
            return self.Trans_Procedure_element(Tx_sig, H_raw, V, D, U, snr)
//  [python]
def Sim_Quan(self, Tx_sig, H_raw=None, H_quan=None, snr=20, Group_enable=True):
    
        H_raw = H_raw[0] + 1j*H_raw[1]
        H_raw = torch.tensor(H_raw).to(self.device)
        H_quan = H_quan[0] + 1j*H_quan[1]
        H_quan = torch.tensor(H_quan).to(self.device)

        V,_,_ = SVD_Precoding_tensor(H_quan, self.P, self.d)
        U,S,V = torch.svd(H_raw)
        D = torch.diag(S)
        if Group_enable:
            return self.Trans_Procedure_group(Tx_sig, H_raw, V, D, U, snr)
        else:
            return self.Trans_Procedure_element(Tx_sig, H_raw, V, D, U, snr)
//  [python]
def Sim_Quan(self, Tx_sig, H_raw=None, H_quan=None, snr=20, Group_enable=True):
    
        H_raw = H_raw[0] + 1j*H_raw[1]
        H_raw = torch.tensor(H_raw).to(self.device)
        H_quan = H_quan[0] + 1j*H_quan[1]
        H_quan = torch.tensor(H_quan).to(self.device)

        V,_,_ = SVD_Precoding_tensor(H_quan, self.P, self.d)
        U,D,_ = torch.linalg.svd(H_raw, full_matrices=False)
        if Group_enable:
            return self.Trans_Procedure_group(Tx_sig, H_raw, V, D, U, snr)
        else:
            return self.Trans_Procedure_element(Tx_sig, H_raw, V, D, U, snr)
//  [python]
def _spectral_init(m, dim_ratio: Union[int, float] = 1):
        u, s, vh = torch.svd(m)
        u = u @ torch.diag(torch.sqrt(s))
        vh = torch.diag(torch.sqrt(s)) @ vh
        if dim_ratio < 1:
            dims = int(u.size(1) * dim_ratio)
            u = u[:, :dims]
            vh = vh[:dims, :]
            # s_share = s[:dims].sum() / s.sum() * 100
            # print(f"SVD eigenvalue share {s_share:.2f}%")
        return u, vh
//  [python]
def _spectral_init(m, dim_ratio: Union[int, float] = 1):
        u, s, vh = torch.linalg.svd(m, full_matrices=False)
        u = u @ torch.diag(torch.sqrt(s))
        vh = torch.diag(torch.sqrt(s)) @ vh
        if dim_ratio < 1:
            dims = int(u.size(1) * dim_ratio)
            u = u[:, :dims]
            vh = vh[:dims, :]
            # s_share = s[:dims].sum() / s.sum() * 100
            # print(f"SVD eigenvalue share {s_share:.2f}%")
        return u, vh
//  [python]
def compute_kernel_bias(batch_hidden_states):
    """for final transformation: y = (x + bias).dot(kernel)
    batched_hidden_states .shape = (batch_size, hidden_dim)
    """
    mu = batch_hidden_states.mean(axis=0, keepdims=True)  # (1, hidden_dim)
    cov = torch.cov(batch_hidden_states.t())  # (hidden_dim, hidden_dim)
    u, s, vh = torch.svd(cov)  # u.shape = (hidden_dim, hidden_dim)  s.shape = (hidden_dim)  vh.shape = (hidden_dim, hidden_dim)
    W = torch.mm(u, torch.diag(1/torch.sqrt(s)))  # (hidden_dim, hidden_dim)
    # kernel = W  # (hidden_dim, hidden_dim)
    # bias = -mu  # (batch_size, hidden_dim)
    return W, -mu
//  [python]
def compute_kernel_bias(batch_hidden_states):
    """for final transformation: y = (x + bias).dot(kernel)
    batched_hidden_states .shape = (batch_size, hidden_dim)
    """
    mu = batch_hidden_states.mean(axis=0, keepdims=True)  # (1, hidden_dim)
    cov = torch.cov(batch_hidden_states.t())  # (hidden_dim, hidden_dim)
    u, s, vh = torch.linalg.svd(cov)  # u.shape = (hidden_dim, hidden_dim)  s.shape = (hidden_dim)  vh.shape = (hidden_dim, hidden_dim)
    W = torch.mm(u, torch.diag(1/torch.sqrt(s)))  # (hidden_dim, hidden_dim)
    # kernel = W  # (hidden_dim, hidden_dim)
    # bias = -mu  # (batch_size, hidden_dim)
    return W, -mu
//  [python]
def rigid_transform_Kabsch_3D_torch(
        A: Tensor, # shape (3, N)
        B: Tensor, # ref shape (3, N)
):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    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, dim = 1, keepdim = True)
    centroid_B = torch.mean(B, dim = 1, keepdim = True)

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

    H = Am @ Bm.T

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

    # Compared with SVD_rigid_transform_torch‘s implementation R^T because (N,3_a,3_b) \matmul (N, 3_b) -> (N, 3_b) @ (N, 3_b, 3_a)
    # Here the coordinates (3, N) so we can use R, so (N, 3_a, 3_b) @ (N, 3_b)
    R = Vt.T @ U.T # R \equiv (N, 3_a, 3_b)
    if torch.det(R) < 0:
        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

    t = -R @ centroid_A + centroid_B
    return R, t