jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_PCA_transform(self,X):
        """
        Computes filewise PCA
        given in https://kaldi-asr.org/doc/ivector-plda-scoring-dense_8cc_source.html
        Apply transform on mean shifted xvectors

        Parameters
        ----------
        Xvectors

        Returns
        ----------
        new xvectors and transform

        """
        # xvec = X[0].cpu().detach().numpy() #N X D
        xvec = X[0]
        num_rows = xvec.shape[0]
        num_cols = xvec.shape[1]
        mean = torch.mean(xvec,0,keepdim=True)
        # xvec = xvec - mean
        S = torch.transpose(xvec,0,1) @ xvec
        # S = torch.matmul(xvec.T,xvec)
        # S = np.cov(xvec.T)
        S = S/num_rows
        
        S = S - torch.transpose(mean,0,1) @ mean
        # ev_s , eig_s = np.linalg.eig(S)
        # ev_s, eig_s , _ = np.linalg.svd(S,full_matrices=True)
        
        
        try:
            ev_s,eig_s,_ = torch.svd(S)
        except:
            bp()
        total_energy = torch.sum(eig_s)
        energy =0.0
        dim=1
        while energy/total_energy <= self.target_energy:
            energy += eig_s[dim-1]
            dim +=1
        transform = ev_s[:,:dim]
        # transform[:,1] = -1 *transform[:,1]
        transxvec = xvec @ transform
        newX = transxvec[np.newaxis]
        # newX = torch.from_numpy(transxvec[np.newaxis]).float().to(self.device)
        # transform = torch.from_numpy(ev_s[:,:dim])

        return newX, torch.transpose(transform,0,1)
//  [python]
def compute_PCA_transform(self,X):
        """
        Computes filewise PCA
        given in https://kaldi-asr.org/doc/ivector-plda-scoring-dense_8cc_source.html
        Apply transform on mean shifted xvectors

        Parameters
        ----------
        Xvectors

        Returns
        ----------
        new xvectors and transform

        """
        # xvec = X[0].cpu().detach().numpy() #N X D
        xvec = X[0]
        num_rows = xvec.shape[0]
        num_cols = xvec.shape[1]
        mean = torch.mean(xvec,0,keepdim=True)
        # xvec = xvec - mean
        S = torch.transpose(xvec,0,1) @ xvec
        # S = torch.matmul(xvec.T,xvec)
        # S = np.cov(xvec.T)
        S = S/num_rows
        
        S = S - torch.transpose(mean,0,1) @ mean
        # ev_s , eig_s = np.linalg.eig(S)
        # ev_s, eig_s , _ = np.linalg.svd(S,full_matrices=True)
        
        
        try:
            ev_s,eig_s,_ = torch.linalg.svd(S)
        except:
            bp()
        total_energy = torch.sum(eig_s)
        energy =0.0
        dim=1
        while energy/total_energy <= self.target_energy:
            energy += eig_s[dim-1]
            dim +=1
        transform = ev_s[:,:dim]
        # transform[:,1] = -1 *transform[:,1]
        transxvec = xvec @ transform
        newX = transxvec[np.newaxis]
        # newX = torch.from_numpy(transxvec[np.newaxis]).float().to(self.device)
        # transform = torch.from_numpy(ev_s[:,:dim])

        return newX, torch.transpose(transform,0,1)
//  [python]
def compute_PCA_transform_withtarget(self,X):
        """
        Computes filewise PCA
        given in https://kaldi-asr.org/doc/ivector-plda-scoring-dense_8cc_source.html
        Apply transform on mean shifted xvectors

        Parameters
        ----------
        Xvectors

        Returns
        ----------
        new xvectors and transform

        """
        # xvec = X[0].cpu().detach().numpy() #N X D
        xvec = X[0]
        num_rows = xvec.shape[0]
        num_cols = xvec.shape[1]
        mean = torch.mean(xvec,0,keepdim=True)
      
        S = torch.transpose(xvec,0,1) @ xvec
        S = S/num_rows
        
        S = S - torch.transpose(mean,0,1) @ mean
        # ev_s , eig_s = np.linalg.eig(S)
        # ev_s, eig_s , _ = np.linalg.svd(S,full_matrices=True)
        
        
        try:
            ev_s,eig_s,_ = torch.svd(S)
        except:
            print('SVD_error')
        # dim = self.pca_dimension
        total_energy = torch.sum(eig_s)
        energy =0.0
        dim=1
        while energy/total_energy <= self.target_energy:
            energy += eig_s[dim-1]
            dim +=1
        print('pca_dim computed: ',dim)
        transform = ev_s[:,:dim]
        # transform[:,1] = -1 *transform[:,1]
        transxvec = xvec @ transform
        newX = transxvec[np.newaxis]
        # newX = torch.from_numpy(transxvec[np.newaxis]).float().to(self.device)
        # transform = torch.from_numpy(ev_s[:,:dim])
        return newX, torch.transpose(transform,0,1)
//  [python]
def compute_PCA_transform_withtarget(self,X):
        """
        Computes filewise PCA
        given in https://kaldi-asr.org/doc/ivector-plda-scoring-dense_8cc_source.html
        Apply transform on mean shifted xvectors

        Parameters
        ----------
        Xvectors

        Returns
        ----------
        new xvectors and transform

        """
        # xvec = X[0].cpu().detach().numpy() #N X D
        xvec = X[0]
        num_rows = xvec.shape[0]
        num_cols = xvec.shape[1]
        mean = torch.mean(xvec,0,keepdim=True)
      
        S = torch.transpose(xvec,0,1) @ xvec
        S = S/num_rows
        
        S = S - torch.transpose(mean,0,1) @ mean
        # ev_s , eig_s = np.linalg.eig(S)
        # ev_s, eig_s , _ = np.linalg.svd(S,full_matrices=True)
        
        
        try:
            ev_s,eig_s,_ = torch.linalg.svd(S)
        except:
            print('SVD_error')
        # dim = self.pca_dimension
        total_energy = torch.sum(eig_s)
        energy =0.0
        dim=1
        while energy/total_energy <= self.target_energy:
            energy += eig_s[dim-1]
            dim +=1
        print('pca_dim computed: ',dim)
        transform = ev_s[:,:dim]
        # transform[:,1] = -1 *transform[:,1]
        transxvec = xvec @ transform
        newX = transxvec[np.newaxis]
        # newX = torch.from_numpy(transxvec[np.newaxis]).float().to(self.device)
        # transform = torch.from_numpy(ev_s[:,:dim])
        return newX, torch.transpose(transform,0,1)
//  [python]
def get_jacobian_singular_values(model, data_loader, num_classes, n_loop=20, print_func=print, rand_data=False):
    jsv, jsv_diff, condition_number = [], [], []
    if rand_data:
        picked_batch = np.random.permutation(len(data_loader))[:n_loop]
    else:
        picked_batch = list(range(n_loop))
    for i, (images, target) in enumerate(data_loader):
        if i in picked_batch:
            images, target = images.cuda(), target.cuda()
            batch_size = images.size(0)
            images.requires_grad = True  # for Jacobian computation
            output = model(images)
            jacobian = compute_jacobian(images,
                                        output)  # shape [batch_size, num_classes, num_channels, input_width, input_height]
            jacobian = jacobian.view(batch_size, num_classes,
                                     -1)  # shape [batch_size, num_classes, num_channels*input_width*input_height]
            u, s, v = torch.svd(
                jacobian)  # u: [batch_size, num_channels*input_width*input_height, num_classes], s: [batch_size, num_classes], v: [batch_size, num_channels*input_width*input_height, num_classes]
            s = s.data.cpu().numpy()
            jsv.append(s)
            jsv_diff.append((s - 1) ** 2)
            condition_number.append(s.max(axis=1) / s.min(axis=1))
            print_func('[%3d/%3d] calculating Jacobian...' % (i, len(data_loader)))
    jsv = np.concatenate(jsv)
    condition_number = np.concatenate(condition_number)
    return jsv, jsv_diff, condition_number
//  [python]
def get_jacobian_singular_values(model, data_loader, num_classes, n_loop=20, print_func=print, rand_data=False):
    jsv, jsv_diff, condition_number = [], [], []
    if rand_data:
        picked_batch = np.random.permutation(len(data_loader))[:n_loop]
    else:
        picked_batch = list(range(n_loop))
    for i, (images, target) in enumerate(data_loader):
        if i in picked_batch:
            images, target = images.cuda(), target.cuda()
            batch_size = images.size(0)
            images.requires_grad = True  # for Jacobian computation
            output = model(images)
            jacobian = compute_jacobian(images,
                                        output)  # shape [batch_size, num_classes, num_channels, input_width, input_height]
            jacobian = jacobian.view(batch_size, num_classes,
                                     -1)  # shape [batch_size, num_classes, num_channels*input_width*input_height]
            u, s, v = torch.linalg.svd(
                jacobian)  # u: [batch_size, num_channels*input_width*input_height, num_classes], s: [batch_size, num_classes], v: [batch_size, num_channels*input_width*input_height, num_classes]
            s = s.data.cpu().numpy()
            jsv.append(s)
            jsv_diff.append((s - 1) ** 2)
            condition_number.append(s.max(axis=1) / s.min(axis=1))
            print_func('[%3d/%3d] calculating Jacobian...' % (i, len(data_loader)))
    jsv = np.concatenate(jsv)
    condition_number = np.concatenate(condition_number)
    return jsv, jsv_diff, condition_number
//  [python]
def compute_similarity_transform(S1, S2):
    '''
    A port of MATLAB's `procrustes` function to PyTorch.
    Adapted from http://stackoverflow.com/a/18927641/1884420
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[1] != 3 and S1.shape[1] != 2:
        S1 = S1.permute(0, 2, 1)
        S2 = S2.permute(0, 2, 1)
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = torch.mean(S1, dim=-1, keepdim=True)
    mu2 = torch.mean(S2, dim=-1, keepdim=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1 ** 2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0, 2, 1)) + 1e-8

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, V = torch.svd(K)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0], 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.det((U.bmm(V.permute(0, 2, 1)))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0, 2, 1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0, 2, 1)

    return S1_hat
//  [python]
def compute_similarity_transform(S1, S2):
    '''
    A port of MATLAB's `procrustes` function to PyTorch.
    Adapted from http://stackoverflow.com/a/18927641/1884420
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[1] != 3 and S1.shape[1] != 2:
        S1 = S1.permute(0, 2, 1)
        S2 = S2.permute(0, 2, 1)
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = torch.mean(S1, dim=-1, keepdim=True)
    mu2 = torch.mean(S2, dim=-1, keepdim=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1 ** 2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0, 2, 1)) + 1e-8

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, V = torch.linalg.svd(K)

    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0], 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.det((U.bmm(V.permute(0, 2, 1)))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0, 2, 1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0, 2, 1)

    return S1_hat
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    # U, _, V = torch.svd(S_non_zero)
    device = S_non_zero.device
    U, _, V = torch.svd(S_non_zero.cpu())
    U = U.to(device=device)
    V = V.to(device=device)

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    # rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))
    det_u_v = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v
    rot_mat_non_zero = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))

    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    pelvis_child = [int(children[0])]
    for i in range(1, parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    rest_mat = []
    target_mat = []
    for child in pelvis_child:
        rest_mat.append(rel_rest_pose[:, child].clone())
        target_mat.append(rel_pose_skeleton[:, child].clone())

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(rest_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    mask_zero = S.sum(dim=(1, 2))

    S_non_zero = S[mask_zero != 0].reshape(-1, 3, 3)

    # U, _, V = torch.svd(S_non_zero)
    device = S_non_zero.device
    U, S, V = torch.linalg.svd(S_non_zero.cpu())
    U = U.to(device=device)
    V = V.to(device=device)

    rot_mat = torch.zeros_like(S)
    rot_mat[mask_zero == 0] = torch.eye(3, device=S.device)

    # rot_mat_non_zero = torch.bmm(V, U.transpose(1, 2))
    det_u_v = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v
    rot_mat_non_zero = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))

    rot_mat[mask_zero != 0] = rot_mat_non_zero

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('rot_mat', rot_mat)

    return rot_mat
//  [python]
def batch_get_3children_orient_svd(rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype):
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(
            rot_mat_chain_parent.transpose(1, 2),
            target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    # U, _, V = torch.svd(S)
    device = S.device
    U, _, V = torch.svd(S.cpu())
    U = U.to(device=device)
    V = V.to(device=device)

    # rot_mat = torch.bmm(V, U.transpose(1, 2))
    det_u_v = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v
    rot_mat = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat
//  [python]
def batch_get_3children_orient_svd(rel_pose_skeleton, rel_rest_pose, rot_mat_chain_parent, children_list, dtype):
    rest_mat = []
    target_mat = []
    for c, child in enumerate(children_list):
        if isinstance(rel_pose_skeleton, list):
            target = rel_pose_skeleton[c].clone()
            template = rel_rest_pose[c].clone()
        else:
            target = rel_pose_skeleton[:, child].clone()
            template = rel_rest_pose[:, child].clone()

        target = torch.matmul(
            rot_mat_chain_parent.transpose(1, 2),
            target)

        target_mat.append(target)
        rest_mat.append(template)

    rest_mat = torch.cat(rest_mat, dim=2)
    target_mat = torch.cat(target_mat, dim=2)
    S = rest_mat.bmm(target_mat.transpose(1, 2))

    # U, _, V = torch.svd(S)
    device = S.device
    U, S, V = torch.linalg.svd(S.cpu())
    U = U.to(device=device)
    V = V.to(device=device)

    # rot_mat = torch.bmm(V, U.transpose(1, 2))
    det_u_v = torch.det(torch.bmm(V, U.transpose(1, 2)))
    det_modify_mat = torch.eye(3, device=U.device).unsqueeze(0).expand(U.shape[0], -1, -1).clone()
    det_modify_mat[:, 2, 2] = det_u_v
    rot_mat = torch.bmm(torch.bmm(V, det_modify_mat), U.transpose(1, 2))

    assert torch.sum(torch.isnan(rot_mat)) == 0, ('3children rot_mat', rot_mat)
    return rot_mat
//  [python]
def rigid_transform_3D(A, B):
    batch, n, dim = A.shape
    centroid_A = torch.mean(A, dim = 1)
    centroid_B = torch.mean(B, dim = 1)
    H = torch.matmul((A - centroid_A[:,None]).permute(0,2,1), B - centroid_B[:,None]) / n
    U, s, V = torch.svd(H)
    R = torch.matmul(V, torch.transpose(U, 1,2))

    negative_det = torch.linalg.det(R) < 0
    s[negative_det, -1] = -s[negative_det, -1]
    V[negative_det, :, 2] = -V[negative_det, :, 2]
    R[negative_det] = torch.matmul(V[negative_det], torch.transpose(U[negative_det],1,2))

    varP = torch.var(A, unbiased=False, dim=1).sum(-1)
    c = 1/varP * torch.sum(s, dim=-1) 

    t = -torch.matmul(c[:,None,None]*R, centroid_A[...,None])[...,-1] + centroid_B
    return c, R, t
//  [python]
def rigid_transform_3D(A, B):
    batch, n, dim = A.shape
    centroid_A = torch.mean(A, dim = 1)
    centroid_B = torch.mean(B, dim = 1)
    H = torch.matmul((A - centroid_A[:,None]).permute(0,2,1), B - centroid_B[:,None]) / n
    U, s, V = torch.linalg.svd(H)
    R = torch.matmul(V, torch.transpose(U, 1,2))

    negative_det = torch.linalg.det(R) < 0
    s[negative_det, -1] = -s[negative_det, -1]
    V[negative_det, :, 2] = -V[negative_det, :, 2]
    R[negative_det] = torch.matmul(V[negative_det], torch.transpose(U[negative_det],1,2))

    varP = torch.var(A, unbiased=False, dim=1).sum(-1)
    c = 1/varP * torch.sum(s, dim=-1) 

    t = -torch.matmul(c[:,None,None]*R, centroid_A[...,None])[...,-1] + centroid_B
    return c, R, t
//  [python]
def svd(self, feat, n=1):
        size = feat.shape
        assert len(size) == 4

        x = feat.view(-1, size[1], size[2] * size[2]).transpose(-2, -1)
        u, s, v = torch.svd(x)

        u = self.removenan(u)
        s = self.removenan(s)
        v = self.removenan(v)

        if n > 0:
            u = F.normalize(u[:, :, :n], dim=1)
            s = F.normalize(s[:, :n], dim=1)
            v = F.normalize(v[:, :, :n], dim=1)

        return u, s, v
//  [python]
def svd(self, feat, n=1):
        size = feat.shape
        assert len(size) == 4

        x = feat.view(-1, size[1], size[2] * size[2]).transpose(-2, -1)
        u, s, v = torch.linalg.svd(x)

        u = self.removenan(u)
        s = self.removenan(s)
        v = self.removenan(v)

        if n > 0:
            u = F.normalize(u[:, :, :n], dim=1)
            s = F.normalize(s[:, :n], dim=1)
            v = F.normalize(v[:, :, :n], dim=1)

        return u, s, v
//  [python]
def Evs(input,k):
    """calcualte EV_{k}(h) from https://arxiv.org/pdf/2005.02178.pdf"""
    _,s,_ = torch.svd(input)
    ek = [s[i]*s[i] for i in range(k)]
    ed = [s[i]*s[i] for i in range(len(s)-k,len(s))]
    return sum(ek)/sum(ed)
//  [python]
def Evs(input,k):
    """calcualte EV_{k}(h) from https://arxiv.org/pdf/2005.02178.pdf"""
    U, s, Vh = torch.linalg.svd(input)
    ek = [s[i]*s[i] for i in range(k)]
    ed = [s[i]*s[i] for i in range(len(s)-k,len(s))]
    return sum(ek)/sum(ed)
//  [python]
def svd_decom(self, tau, X):
        U, S, Vh = torch.svd(X, some=True)
        S = torch.diag(S)
        S = self.shrinkage(tau, S)
        return torch.mm(torch.mm(U, S), torch.transpose(Vh, 0, 1))
//  [python]
def svd_decom(self, tau, X):
        U, S, Vh = torch.linalg.svd(X, full_matrices=False)
        S = torch.diag(S)
        S = self.shrinkage(tau, S)
        return torch.mm(torch.mm(U, S), torch.transpose(Vh, 0, 1))
//  [python]
def torch_pca_svd(X, center=True):
    """
    perform PCA over X
    - X is of size (num_points, vec_size)

    NOTE: unlike PCA_eig, the number of vectors/values returned is: min(num_points, vec_size)
    """
    n, _ = X.shape
    # center points along axes
    if center:
        X = X - X.mean(dim=0)
    # perform singular value decomposition
    u, s, v = torch.svd(X)
    # sort components by decreasing variance
    # these are already sorted?
    components = v.T
    explained_variance = torch.mul(s, s) / (n - 1)
    return components, explained_variance
//  [python]
def torch_pca_svd(X, center=True):
    """
    perform PCA over X
    - X is of size (num_points, vec_size)

    NOTE: unlike PCA_eig, the number of vectors/values returned is: min(num_points, vec_size)
    """
    n, _ = X.shape
    # center points along axes
    if center:
        X = X - X.mean(dim=0)
    # perform singular value decomposition
    u, s, v = torch.linalg.svd(X)
    # sort components by decreasing variance
    # these are already sorted?
    components = v.T
    explained_variance = torch.mul(s, s) / (n - 1)
    return components, explained_variance
//  [python]
def get_vne(H):
    Z = torch.nn.functional.normalize(H, dim=1)
    sing_val = torch.svd(Z / np.sqrt(Z.shape[0]))[1]
    eig_val = sing_val ** 2
    return - (eig_val * torch.log(eig_val)).nansum()
//  [python]
def get_vne(H):
    Z = torch.nn.functional.normalize(H, dim=1)
    sing_val = torch.linalg.svd(Z / np.sqrt(Z.shape[0]))[1]
    eig_val = sing_val ** 2
    return - (eig_val * torch.log(eig_val)).nansum()
//  [python]
def __init__(self,encode_layers,bilinear_layers,Nkoopman,u_dim):
        super(Network,self).__init__()
        ELayers = OrderedDict()
        for layer_i in range(len(encode_layers)-1):
            ELayers["linear_{}".format(layer_i)] = nn.Linear(encode_layers[layer_i],encode_layers[layer_i+1])
            if layer_i != len(encode_layers)-2:
                ELayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.encode_net = nn.Sequential(ELayers)
        BLayers = OrderedDict()
        for layer_i in range(len(bilinear_layers)-1):
            BLayers["linear_{}".format(layer_i)] = nn.Linear(bilinear_layers[layer_i],bilinear_layers[layer_i+1])
            if layer_i != len(bilinear_layers)-2:
                BLayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.bilinear_net = nn.Sequential(BLayers)         
        self.Nkoopman = Nkoopman
        self.u_dim = u_dim
        self.lA = nn.Linear(Nkoopman,Nkoopman,bias=False)
        self.lA.weight.data = gaussian_init_(Nkoopman, std=1)
        U, _, V = torch.svd(self.lA.weight.data)
        self.lA.weight.data = torch.mm(U, V.t()) * 0.9
        self.lB = nn.Linear(bilinear_layers[-1],Nkoopman,bias=False)
//  [python]
def __init__(self,encode_layers,bilinear_layers,Nkoopman,u_dim):
        super(Network,self).__init__()
        ELayers = OrderedDict()
        for layer_i in range(len(encode_layers)-1):
            ELayers["linear_{}".format(layer_i)] = nn.Linear(encode_layers[layer_i],encode_layers[layer_i+1])
            if layer_i != len(encode_layers)-2:
                ELayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.encode_net = nn.Sequential(ELayers)
        BLayers = OrderedDict()
        for layer_i in range(len(bilinear_layers)-1):
            BLayers["linear_{}".format(layer_i)] = nn.Linear(bilinear_layers[layer_i],bilinear_layers[layer_i+1])
            if layer_i != len(bilinear_layers)-2:
                BLayers["relu_{}".format(layer_i)] = nn.ReLU()
        self.bilinear_net = nn.Sequential(BLayers)         
        self.Nkoopman = Nkoopman
        self.u_dim = u_dim
        self.lA = nn.Linear(Nkoopman,Nkoopman,bias=False)
        self.lA.weight.data = gaussian_init_(Nkoopman, std=1)
        U, S, V = torch.linalg.svd(self.lA.weight.data)
        self.lA.weight.data = torch.mm(U, V.t()) * 0.9
        self.lB = nn.Linear(bilinear_layers[-1],Nkoopman,bias=False)
//  [python]
def rigid_transform_2d(self, A, B, weights=None, weight_threshold=0):
        """
        Input:
            - A:       [bs, num_corr, 2], source point cloud
            - B:       [bs, num_corr, 2], target point cloud
            - weights: [bs, num_corr]     weight for each correspondence
            - weight_threshold: float,    clips points with weight below threshold
        Output:
            - R, t
        """
        bs = A.shape[0]
        if weights is None:
            weights = torch.ones_like(A[:, :, 0])
        weights[weights < weight_threshold] = 0

        # find mean of point cloud
        centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (
                torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
        centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (
                torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
        # subtract mean
        Am = A - centroid_A
        Bm = B - centroid_B

        # construct weight covariance matrix
        Weight = torch.diag_embed(weights)
        H = Am.permute(0, 2, 1) @ Weight @ Bm
        U, S, Vt = torch.svd(H.cpu())
        U, S, Vt = U.to(weights.device), S.to(weights.device), Vt.to(weights.device)

        delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
        eye = torch.eye(2)[None, :, :].repeat(bs, 1, 1).to(A.device)
        eye[:, -1, -1] = delta_UV
        R = Vt @ eye @ U.permute(0, 2, 1)
        t = centroid_B.permute(0, 2, 1) - R @ centroid_A.permute(0, 2, 1)
        # print('Estimated R:\n', R)
        # print('Estimated T:\n', t)
        return self.integrate_trans(R, t)
//  [python]
def rigid_transform_2d(self, A, B, weights=None, weight_threshold=0):
        """
        Input:
            - A:       [bs, num_corr, 2], source point cloud
            - B:       [bs, num_corr, 2], target point cloud
            - weights: [bs, num_corr]     weight for each correspondence
            - weight_threshold: float,    clips points with weight below threshold
        Output:
            - R, t
        """
        bs = A.shape[0]
        if weights is None:
            weights = torch.ones_like(A[:, :, 0])
        weights[weights < weight_threshold] = 0

        # find mean of point cloud
        centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (
                torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
        centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (
                torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-6)
        # subtract mean
        Am = A - centroid_A
        Bm = B - centroid_B

        # construct weight covariance matrix
        Weight = torch.diag_embed(weights)
        H = Am.permute(0, 2, 1) @ Weight @ Bm
        U, S, Vt = torch.linalg.svd(H.cpu())
        U, S, Vt = U.to(weights.device), S.to(weights.device), Vt.to(weights.device)

        delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
        eye = torch.eye(2)[None, :, :].repeat(bs, 1, 1).to(A.device)
        eye[:, -1, -1] = delta_UV
        R = Vt @ eye @ U.permute(0, 2, 1)
        t = centroid_B.permute(0, 2, 1) - R @ centroid_A.permute(0, 2, 1)
        # print('Estimated R:\n', R)
        # print('Estimated T:\n', t)
        return self.integrate_trans(R, t)
//  [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_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, V = torch.svd(B_t)
        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, V = torch.svd(B_t)
        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_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, V = torch.linalg.svd(B_t)
        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, V = torch.linalg.svd(B_t)
        U = Q.matmul(U)

    return U, S, V
//  [python]
def loss_byol_directpred_grad(self, z1, z2, z1m, z2m):
        # calculate F & wp
        tmp_F = (torch.mm(z1m.t(), z1m) + torch.mm(z2m.t(), z2m)) / (2*z1m.shape[0])
        torch.distributed.all_reduce(tmp_F)
        tmp_F = tmp_F / torch.distributed.get_world_size()
        if self.F is None:
            self.F = tmp_F.detach()
        else:
            self.F = self.cfg.rho * self.F + (1 - self.cfg.rho) * tmp_F.detach()
        u, s, v = torch.svd(self.F.float(), compute_uv=True)
        max_s = torch.max(s)
        sqrt_F = torch.mm(u, torch.mm(torch.diag(torch.sqrt(s)), v.t())).half()
        s = s / torch.max(s)
        s = torch.sqrt(s) + self.cfg.eps
        s = s.clamp(1e-4)
        wp = torch.mm(u, torch.mm(torch.diag(s), v.t())).half()
        eps_mat = self.cfg.eps**2*torch.eye(z1.shape[1]).cuda()

        # loss1 is made up of z1 & z2m
        # calculate common denominator
        common_denominator = torch.linalg.norm(torch.mm(z1, wp), dim=-1).unsqueeze(-1) # (n, 1)

        # calculate positive term
        pos_term = - torch.mm(z2m, wp) # (n, c)

        # calculate lambda
        neg_denominator = torch.einsum('nc,nc->n', [torch.mm(z1, self.F/max_s + eps_mat), z1]).unsqueeze(-1) # (n, 1)
        neg_numerator = torch.einsum('nc,nc->n', [torch.mm(z1, wp), z2m]).unsqueeze(-1) # (n, 1)
        neg_lambda = neg_numerator / neg_denominator

        # calculate negative term
        neg_term = torch.mm(z1, self.F/max_s + eps_mat) # (n, c)

        grad1 = 1/common_denominator * (pos_term + neg_lambda * neg_term) # (n, c)
        loss1 = (z1 * grad1.detach()).sum(-1).mean()

        # loss2 is made up of z2 & z1m
        # calculate common denominator
        common_denominator = torch.linalg.norm(torch.mm(z2, wp), dim=-1).unsqueeze(-1) # (n, 1)

        # calculate positive term
        pos_term = - torch.mm(z1m, wp) # (n, c)

        # calculate lambda
        neg_denominator = torch.einsum('nc,nc->n', [torch.mm(z2, self.F/max_s + eps_mat), z2]).unsqueeze(-1) # (n, 1)
        neg_numerator = torch.einsum('nc,nc->n', [torch.mm(z2, wp), z1m]).unsqueeze(-1) # (n, 1)
        neg_lambda = neg_numerator / neg_denominator

        # calculate first negative term
        neg_term = torch.mm(z2, self.F/max_s + eps_mat) # (n, c)

        grad2 = 1/common_denominator * (pos_term + neg_lambda * neg_term) # (n, c)
        loss2 = (z2 * grad2.detach()).sum(-1).mean()
        loss = 0.5 * (loss1 + loss2)

        # compute positive similarity, just for observation
        pos_sim1 = torch.einsum('nc,nc->n', [z1, z2m]).mean().detach()
        pos_sim2 = torch.einsum('nc,nc->n', [z2, z1m]).mean().detach()
        pos_sim = 0.5 * (pos_sim1 + pos_sim2)
        
        return loss, pos_sim
//  [python]
def loss_byol_directpred_grad(self, z1, z2, z1m, z2m):
        # calculate F & wp
        tmp_F = (torch.mm(z1m.t(), z1m) + torch.mm(z2m.t(), z2m)) / (2*z1m.shape[0])
        torch.distributed.all_reduce(tmp_F)
        tmp_F = tmp_F / torch.distributed.get_world_size()
        if self.F is None:
            self.F = tmp_F.detach()
        else:
            self.F = self.cfg.rho * self.F + (1 - self.cfg.rho) * tmp_F.detach()
        u, s, v = torch.linalg.svd(self.F.float())
        max_s = torch.max(s)
        sqrt_F = torch.mm(u, torch.mm(torch.diag(torch.sqrt(s)), v.t())).half()
        s = s / torch.max(s)
        s = torch.sqrt(s) + self.cfg.eps
        s = s.clamp(1e-4)
        wp = torch.mm(u, torch.mm(torch.diag(s), v.t())).half()
        eps_mat = self.cfg.eps**2*torch.eye(z1.shape[1]).cuda()

        # loss1 is made up of z1 & z2m
        # calculate common denominator
        common_denominator = torch.linalg.norm(torch.mm(z1, wp), dim=-1).unsqueeze(-1) # (n, 1)

        # calculate positive term
        pos_term = - torch.mm(z2m, wp) # (n, c)

        # calculate lambda
        neg_denominator = torch.einsum('nc,nc->n', [torch.mm(z1, self.F/max_s + eps_mat), z1]).unsqueeze(-1) # (n, 1)
        neg_numerator = torch.einsum('nc,nc->n', [torch.mm(z1, wp), z2m]).unsqueeze(-1) # (n, 1)
        neg_lambda = neg_numerator / neg_denominator

        # calculate negative term
        neg_term = torch.mm(z1, self.F/max_s + eps_mat) # (n, c)

        grad1 = 1/common_denominator * (pos_term + neg_lambda * neg_term) # (n, c)
        loss1 = (z1 * grad1.detach()).sum(-1).mean()

        # loss2 is made up of z2 & z1m
        # calculate common denominator
        common_denominator = torch.linalg.norm(torch.mm(z2, wp), dim=-1).unsqueeze(-1) # (n, 1)

        # calculate positive term
        pos_term = - torch.mm(z1m, wp) # (n, c)

        # calculate lambda
        neg_denominator = torch.einsum('nc,nc->n', [torch.mm(z2, self.F/max_s + eps_mat), z2]).unsqueeze(-1) # (n, 1)
        neg_numerator = torch.einsum('nc,nc->n', [torch.mm(z2, wp), z1m]).unsqueeze(-1) # (n, 1)
        neg_lambda = neg_numerator / neg_denominator

        # calculate first negative term
        neg_term = torch.mm(z2, self.F/max_s + eps_mat) # (n, c)

        grad2 = 1/common_denominator * (pos_term + neg_lambda * neg_term) # (n, c)
        loss2 = (z2 * grad2.detach()).sum(-1).mean()
        loss = 0.5 * (loss1 + loss2)

        # compute positive similarity, just for observation
        pos_sim1 = torch.einsum('nc,nc->n', [z1, z2m]).mean().detach()
        pos_sim2 = torch.einsum('nc,nc->n', [z2, z1m]).mean().detach()
        pos_sim = 0.5 * (pos_sim1 + pos_sim2)
        
        return loss, pos_sim
//  [python]
def reset_parameters(self) -> None:
        temp = torch.empty(
            self.grid_dim_y * self.miniblock, self.grid_dim_x * self.miniblock, device=self.device
        )
        init.kaiming_normal_(temp)
        temp = temp.view(self.grid_dim_y, self.miniblock, self.grid_dim_x, self.miniblock).permute(0, 2, 1, 3)
        _, s, _ = torch.svd(temp, compute_uv=False)
        self.weight.data.copy_(s)

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def reset_parameters(self) -> None:
        temp = torch.empty(
            self.grid_dim_y * self.miniblock, self.grid_dim_x * self.miniblock, device=self.device
        )
        init.kaiming_normal_(temp)
        temp = temp.view(self.grid_dim_y, self.miniblock, self.grid_dim_x, self.miniblock).permute(0, 2, 1, 3)
        _, s, _ = torch.linalg.svd(temp, compute_uv=False)
        self.weight.data.copy_(s)

        if self.bias is not None:
            init.uniform_(self.bias, 0, 0)
//  [python]
def random_rotation(batch_size, dim, device=None):
        """
        Generates a batch of random `dim`-dimensional rotation matrices.
        """
        if dim == 3:
            R = rotation_conversions.random_rotations(batch_size, device=device)
        else:
            # generate random rotation matrices with orthogonalization of
            # random normal square matrices, followed by a transformation
            # that ensures determinant(R)==1
            H = torch.randn(batch_size, dim, dim, dtype=torch.float32, device=device)
            U, _, V = torch.svd(H)
            E = torch.eye(dim, dtype=torch.float32, device=device)[None].repeat(
                batch_size, 1, 1
            )
            E[:, -1, -1] = torch.det(torch.bmm(U, V.transpose(2, 1)))
            R = torch.bmm(torch.bmm(U, E), V.transpose(2, 1))
            assert torch.allclose(torch.det(R), R.new_ones(batch_size), atol=1e-4)

        return R
//  [python]
def random_rotation(batch_size, dim, device=None):
        """
        Generates a batch of random `dim`-dimensional rotation matrices.
        """
        if dim == 3:
            R = rotation_conversions.random_rotations(batch_size, device=device)
        else:
            # generate random rotation matrices with orthogonalization of
            # random normal square matrices, followed by a transformation
            # that ensures determinant(R)==1
            H = torch.randn(batch_size, dim, dim, dtype=torch.float32, device=device)
            U, S, V = torch.linalg.svd(H)
            E = torch.eye(dim, dtype=torch.float32, device=device)[None].repeat(
                batch_size, 1, 1
            )
            E[:, -1, -1] = torch.det(torch.bmm(U, V.transpose(2, 1)))
            R = torch.bmm(torch.bmm(U, E), V.transpose(2, 1))
            assert torch.allclose(torch.det(R), R.new_ones(batch_size), atol=1e-4)

        return R
//  [python]
def kabsch_torch(X, Y):
    """ Kabsch alignment of X into Y. 
        Assumes X,Y are both (D, N) - usually (3, N)
    """
    #  center X and Y to the origin
    X_ = X - X.mean(dim=-1, keepdim=True)
    Y_ = Y - Y.mean(dim=-1, keepdim=True)
    # calculate convariance matrix (for each prot in the batch)
    C = torch.matmul(X_, Y_.t())
    # Optimal rotation matrix via SVD - warning! W must be transposed
    if int(torch.__version__.split(".")[1]) < 8:
        V, S, W = torch.svd(C.detach())
        W = W.t()
    else: 
        V, S, W = torch.linalg.svd(C.detach()) 
    # determinant sign for direction correction
    d = (torch.det(V) * torch.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = torch.matmul(V, W)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    return X_, Y_
//  [python]
def kabsch_torch(X, Y):
    """ Kabsch alignment of X into Y. 
        Assumes X,Y are both (D, N) - usually (3, N)
    """
    #  center X and Y to the origin
    X_ = X - X.mean(dim=-1, keepdim=True)
    Y_ = Y - Y.mean(dim=-1, keepdim=True)
    # calculate convariance matrix (for each prot in the batch)
    C = torch.matmul(X_, Y_.t())
    # Optimal rotation matrix via SVD - warning! W must be transposed
    V, S, W = torch.linalg.svd(C.detach())
    # determinant sign for direction correction
    d = (torch.det(V) * torch.det(W)) < 0.0
    if d:
        S[-1]    = S[-1] * (-1)
        V[:, -1] = V[:, -1] * (-1)
    # Create Rotation matrix U
    U = torch.matmul(V, W)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    return X_, Y_
//  [python]
def whiten_and_color(self, cf, sf):
        # content size
        c_size = cf.size()
        # channels x (height x width)
        # content mean
        c_mean = torch.mean(cf,1)
        c_mean = c_mean.unsqueeze(1).expand_as(cf)
        cf = cf - c_mean
        
        # content conv
        c_conv = torch.mm(cf, cf.t()).div(c_size[1]-1) + torch.eye(c_size[0]).double()
        #c_u, c_e, c_v = torch.linalg.svd(c_conv,some=False)
        c_u, c_e, c_v = torch.svd(c_conv, some=False)

        k_c = c_size[0]
        for i in range(c_size[0]):
            if c_e[i] < 0.00001:
                k_c = i
                break
        

        # style size
        s_size = sf.size()
        # style mean
        s_mean = torch.mean(sf,1)
        sf = sf - s_mean.unsqueeze(1).expand_as(sf)

        # style conv
        s_conv = torch.mm(sf, sf.t()).div(s_size[1]-1)
        s_u, s_e, s_v = torch.svd(s_conv, some=False) #torch.linalg.svd(s_conv, some=False)

        k_s = s_size[0]
        for i in range(s_size[0]):
            if s_e[i] < 0.00001:
                k_s = i
                break


        c_d = (c_e[0:k_c]).pow(-0.5)
        step1 = torch.mm(c_v[:,0:k_c], torch.diag(c_d))
        step2 = torch.mm(step1, (c_v[:,0:k_c].t()))
        whiten_cf = torch.mm(step2, cf)

        s_d = (s_e[0:k_s]).pow(0.5)
        target = torch.mm(torch.mm(torch.mm(s_v[:,0:k_s], torch.diag(s_d)), (s_v[:,0:k_s].t())), whiten_cf)
        target = target + s_mean.unsqueeze(1).expand_as(target)

        return target
//  [python]
def whiten_and_color(self, cf, sf):
        # content size
        c_size = cf.size()
        # channels x (height x width)
        # content mean
        c_mean = torch.mean(cf,1)
        c_mean = c_mean.unsqueeze(1).expand_as(cf)
        cf = cf - c_mean
        
        # content conv
        c_conv = torch.mm(cf, cf.t()).div(c_size[1]-1) + torch.eye(c_size[0]).double()
        c_u, c_e, c_v = torch.linalg.svd(c_conv)

        k_c = c_size[0]
        for i in range(c_size[0]):
            if c_e[i] < 0.00001:
                k_c = i
                break
        

        # style size
        s_size = sf.size()
        # style mean
        s_mean = torch.mean(sf,1)
        sf = sf - s_mean.unsqueeze(1).expand_as(sf)

        # style conv
        s_conv = torch.mm(sf, sf.t()).div(s_size[1]-1)
        s_u, s_e, s_v = torch.linalg.svd(s_conv)

        k_s = s_size[0]
        for i in range(s_size[0]):
            if s_e[i] < 0.00001:
                k_s = i
                break


        c_d = (c_e[0:k_c]).pow(-0.5)
        step1 = torch.mm(c_v[:,0:k_c], torch.diag(c_d))
        step2 = torch.mm(step1, (c_v[:,0:k_c].t()))
        whiten_cf = torch.mm(step2, cf)

        s_d = (s_e[0:k_s]).pow(0.5)
        target = torch.mm(torch.mm(torch.mm(s_v[:,0:k_s], torch.diag(s_d)), (s_v[:,0:k_s].t())), whiten_cf)
        target = target + s_mean.unsqueeze(1).expand_as(target)

        return target
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    """
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    """
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0, 2, 1)
        S2 = S2.permute(0, 2, 1)
        transposed = True
    assert S2.shape[1] == S1.shape[1]
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2
    var1 = torch.sum(X1 ** 2, dim=1).sum(dim=1)
    K = X1.bmm(X2.permute(0, 2, 1))
    U, s, V = torch.svd(K)
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0], 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0, 2, 1))))
    R = V.bmm(Z.bmm(U.permute(0, 2, 1)))
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1
    t = mu2 - scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(mu1)
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t
    if transposed:
        S1_hat = S1_hat.permute(0, 2, 1)
    return S1_hat, (scale, R, t)
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    """
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    """
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0, 2, 1)
        S2 = S2.permute(0, 2, 1)
        transposed = True
    assert S2.shape[1] == S1.shape[1]
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2
    var1 = torch.sum(X1 ** 2, dim=1).sum(dim=1)
    K = X1.bmm(X2.permute(0, 2, 1))
    U, s, V = torch.linalg.svd(K)
    Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0], 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0, 2, 1))))
    R = V.bmm(Z.bmm(U.permute(0, 2, 1)))
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1
    t = mu2 - scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(mu1)
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t
    if transposed:
        S1_hat = S1_hat.permute(0, 2, 1)
    return S1_hat, (scale, R, t)
//  [python]
def __call__(self, data: Data) -> Data:
        if data.x.size(-1) > self.out_channels:
            U, S, _ = torch.svd(data.x)
            data.x = torch.mm(U[:, :self.out_channels],
                              torch.diag(S[:self.out_channels]))
        return data
//  [python]
def __call__(self, data: Data) -> Data:
        if data.x.size(-1) > self.out_channels:
            U, S, _ = torch.linalg.svd(data.x)
            data.x = torch.mm(U[:, :self.out_channels],
                              torch.diag(S[:self.out_channels]))
        return data