jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_R_similarity_matrix(self, point_name, states, idxs, sim_func=generalized_squared_cosine):
        """
        Compute R-based similarity matrix
        :param point_name: Observation point
        :param states: List of states (or list of list) used to compute similarities
        :param idxs: Sample indices in state collection
        :param sim_func: Similarity function (default: generalized_squared_cosine)
        :return: Similarity matrix as torch Tensor
        """
        # Cell states
        cell_states = list()

        # For each states
        for state in states:
            # Get observed data
            cell_states.append(self._get_observer_data(point_name, state, idxs, concat=True, concat_dim=2))
        # end for

        # Similarity matrix
        sim_matrix = torch.zeros(len(cell_states), len(cell_states))

        # For each pair
        for i in range(len(cell_states)):
            for j in range(len(cell_states)):
                # R matrices
                R_i = torch.mm(cell_states[i].t(), cell_states[i])
                R_j = torch.mm(cell_states[j].t(), cell_states[j])

                # SVDs
                Ui, Si, _ = torch.linalg.svd(R_i)
                Uj, Sj, _ = torch.linalg.svd(R_j)

                # Similarity matrix
                sim_matrix[i, j] = sim_func(Si, Ui, Sj, Uj)
            # end for
        # end for

        return sim_matrix
//  [python]
def symmetric_orthogonalization(x):
    """Maps 9D input vectors onto SO(3) via symmetric orthogonalization.

    x: should have size [batch_size, 9]

    Output has size [batch_size, 3, 3], where each inner 3x3 matrix is in SO(3).
    """
    m = x.view(-1, 3, 3)
    u, s, v = torch.svd(m)
    vt = torch.transpose(v, 1, 2)
    det = torch.det(torch.matmul(u, vt))
    det = det.view(-1, 1, 1)
    vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)
    r = torch.matmul(u, vt)
    return r
//  [python]
def symmetric_orthogonalization(x):
    """Maps 9D input vectors onto SO(3) via symmetric orthogonalization.

    x: should have size [batch_size, 9]

    Output has size [batch_size, 3, 3], where each inner 3x3 matrix is in SO(3).
    """
    m = x.view(-1, 3, 3)
    u, s, v = torch.linalg.svd(m)
    vt = torch.transpose(v, 1, 2)
    det = torch.det(torch.matmul(u, vt))
    det = det.view(-1, 1, 1)
    vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)
    r = torch.matmul(u, vt)
    return r
//  [python]
def _F_from_XY(X, Y, W=None, normalize=True, show_debug=False): # Ref: https://github.com/marktao99/python/blob/master/CVP/samples/sfm.py#L55
    # X, Y: [N, 2]
    if normalize:
        # print(X.t().numpy())
        X, Y, T1, T2 = _normalize_XY(X, Y)

    xx = torch.cat([X.t(), Y.t()], dim=0)
    # print(xx.size())
    # print(xx.size())
    XX = torch.stack([
        xx[2, :] * xx[0, :], xx[2, :] * xx[1, :], xx[2, :],
        xx[3, :] * xx[0, :], xx[3, :] * xx[1, :], xx[3, :],
        xx[0, :], xx[1, :], torch.ones_like(xx[0, :])
    ], dim=0).t()
    if W is not None:
        XX = torch.mm(W, XX)
    U, D, V = torch.svd(XX, some=True)
    # print(XX.size(), U.size(), D.size(), V.size())
    # print(V.size(), V[:, -1].reshape([3, 3]).numpy())
    # print(D.numpy())

    # u, d, v = batch_svd(XX.cuda().float().unsqueeze(0))
    # print(v.size(), v[0, :, -1].reshape([3, 3]).cpu().numpy())
    # print(d.cpu().numpy().flatten())

    if show_debug:
        print('[info.Debug@_F_from_XY] Singualr values of XX:\n', D.numpy())
    # print(D.numpy())
    # print(V.numpy().T, V.numpy().shape)
    # print(U.size(), D.size(), V.size(), X.size())
    # print(V[:, -1].numpy()))

    # U_np, D_np, V_np = np.linalg.svd(XX.numpy())
    # V_np = torch.from_numpy(V_np)
    # print(V[-1].numpy())

    F_recover = torch.reshape(V[:, -1], (3, 3))
    # F_recover = torch.reshape(v[0, :, -1].cpu(), (3, 3))

    # return F_recover, np.reshape(V_np[-1], (3, 3))

    # F_recover_rescale = F_recover / torch.norm(F_recover) * torch.norm(F)
    # print('-', F_recover_rescale.numpy())
    FU, FD, FV= torch.svd(F_recover, some=True)
    FDnew = torch.diag(FD)
    FDnew[2, 2] = 0
    # print(FU.size(), FDnew.size(), FV.t().size())
    F_recover_sing = torch.mm(FU, torch.mm(FDnew, FV.t()))
    # F_recover_sing_rescale = F_recover_sing / torch.norm(F_recover_sing) * torch.norm(F)

    if normalize:
        F_recover_sing = torch.mm(T2.t(), torch.mm(F_recover_sing, T1))
    return F_recover_sing
//  [python]
def _F_from_XY(X, Y, W=None, normalize=True, show_debug=False): # Ref: https://github.com/marktao99/python/blob/master/CVP/samples/sfm.py#L55
    # X, Y: [N, 2]
    if normalize:
        # print(X.t().numpy())
        X, Y, T1, T2 = _normalize_XY(X, Y)

    xx = torch.cat([X.t(), Y.t()], dim=0)
    # print(xx.size())
    # print(xx.size())
    XX = torch.stack([
        xx[2, :] * xx[0, :], xx[2, :] * xx[1, :], xx[2, :],
        xx[3, :] * xx[0, :], xx[3, :] * xx[1, :], xx[3, :],
        xx[0, :], xx[1, :], torch.ones_like(xx[0, :])
    ], dim=0).t()
    if W is not None:
        XX = torch.mm(W, XX)
    U, S, V = torch.linalg.svd(XX)
    # print(XX.size(), U.size(), D.size(), V.size())
    # print(V.size(), V[:, -1].reshape([3, 3]).numpy())
    # print(D.numpy())

    # u, d, v = batch_svd(XX.cuda().float().unsqueeze(0))
    # print(v.size(), v[0, :, -1].reshape([3, 3]).cpu().numpy())
    # print(d.cpu().numpy().flatten())

    if show_debug:
        print('[info.Debug@_F_from_XY] Singualr values of XX:\n', S.numpy())
    # print(D.numpy())
    # print(V.numpy().T, V.numpy().shape)
    # print(U.size(), D.size(), V.size(), X.size())
    # print(V[:, -1].numpy()))

    # U_np, D_np, V_np = np.linalg.svd(XX.numpy())
    # V_np = torch.from_numpy(V_np)
    # print(V[-1].numpy())

    F_recover = torch.reshape(V[:, -1], (3, 3))
    # F_recover = torch.reshape(v[0, :, -1].cpu(), (3, 3))

    # return F_recover, np.reshape(V_np[-1], (3, 3))

    # F_recover_rescale = F_recover / torch.norm(F_recover) * torch.norm(F)
    # print('-', F_recover_rescale.numpy())
    FU, FD, FV= torch.linalg.svd(F_recover)
    FDnew = torch.diag(FD)
    FDnew[2, 2] = 0
    # print(FU.size(), FDnew.size(), FV.t().size())
    F_recover_sing = torch.mm(FU, torch.mm(FDnew, FV.t()))
    # F_recover_sing_rescale = F_recover_sing / torch.norm(F_recover_sing) * torch.norm(F)

    if normalize:
        F_recover_sing = torch.mm(T2.t(), torch.mm(F_recover_sing, T1))
    return F_recover_sing
//  [python]
def rigid_transform_torch(self, A, B, M):
    
        assert A.size() == B.size()
    
        batch_size, num_rows, num_cols = A.size()
       
        ## mask to 0/1 weights for motive/static points
        W=M.type(torch.bool).unsqueeze(2)

        # find mean column wise
        centroid_A = torch.mean(A.transpose(2,1).contiguous()*W, axis=1)
        centroid_B = torch.mean(B.transpose(2,1).contiguous()*W, axis=1)
    
        # ensure centroids are 3x1
        centroid_A = centroid_A.reshape(batch_size,num_rows,1)
        centroid_B = centroid_B.reshape(batch_size,num_rows,1)
    
        # subtract mean
        Am = A - centroid_A
        Bm = B - centroid_B
    
        H = torch.matmul(Am, Bm.transpose(2,1).contiguous()*W)

        # find rotation
        U, S, V = torch.svd(H)
        Z = torch.matmul(V,U.transpose(2,1).contiguous())
        # special reflection case
        d= (torch.linalg.det(Z) < 0).type(torch.int8)
        # -1/1 
        d=d*2-1
        Vc = V.clone()
        Vc[:,2,:]*=-d.view(batch_size,1)
        R = torch.matmul(Vc,U.transpose(2,1).contiguous())
       
        t = torch.matmul(-R, centroid_A)+centroid_B
        
        Trans=torch.cat((torch.cat((R,t),axis=2), torch.tensor([0,0,0,1]).repeat(batch_size,1).cuda().view(batch_size,1,4)),axis=1)
    
        return Trans
//  [python]
def rigid_transform_torch(self, A, B, M):
    
        assert A.size() == B.size()
    
        batch_size, num_rows, num_cols = A.size()
       
        ## mask to 0/1 weights for motive/static points
        W=M.type(torch.bool).unsqueeze(2)

        # find mean column wise
        centroid_A = torch.mean(A.transpose(2,1).contiguous()*W, axis=1)
        centroid_B = torch.mean(B.transpose(2,1).contiguous()*W, axis=1)
    
        # ensure centroids are 3x1
        centroid_A = centroid_A.reshape(batch_size,num_rows,1)
        centroid_B = centroid_B.reshape(batch_size,num_rows,1)
    
        # subtract mean
        Am = A - centroid_A
        Bm = B - centroid_B
    
        H = torch.matmul(Am, Bm.transpose(2,1).contiguous()*W)

        # find rotation
        U, S, V = torch.linalg.svd(H)
        Z = torch.matmul(V,U.transpose(2,1).contiguous())
        # special reflection case
        d= (torch.linalg.det(Z) < 0).type(torch.int8)
        # -1/1 
        d=d*2-1
        Vc = V.clone()
        Vc[:,2,:]*=-d.view(batch_size,1)
        R = torch.matmul(Vc,U.transpose(2,1).contiguous())
       
        t = torch.matmul(-R, centroid_A)+centroid_B
        
        Trans=torch.cat((torch.cat((R,t),axis=2), torch.tensor([0,0,0,1]).repeat(batch_size,1).cuda().view(batch_size,1,4)),axis=1)
    
        return Trans
//  [python]
def plot_singular_values(self, point_name, states, idxs, color, linewidth=1, length=-1, show_title=True, title="",
                                   xticks=None, yticks=None, ylim=None, xlim=None, log10=False):
        """
        Plot singular values
        :param point_name: Name of the observation point
        :param states: Observer's state (or None if all)
        :param idxs: Samples indices in the state collection (or None if all)
        :param color: Line color
        :param linewidth: Line width
        :param length: Length of the plot
        :param show_title: Add title to the plot ?
        :param title: Plot title
        :param xticks: X ticks
        :param yticks: Y ticks
        :param ylim: Y limits [min, max]
        :param xlim: X limits [min, max]
        :param log10: If true, show log10 of singular values
        """
        # Get observed data
        cell_states = self._get_observer_data(point_name, states, idxs, concat=True, concat_dim=2)

        # Correlation matrix R
        R = torch.mm(cell_states.t(), cell_states) / cell_states.size(0)

        # SVD on state
        _, S, _ = torch.svd(R)

        # Log10?
        if log10:
            S = torch.log10(S)
        # end if

        # Learning PC energy
        if length != -1:
            plt.plot(S[:length], color=color, linewidth=linewidth)
        else:
            plt.plot(S, color=color, linewidth=linewidth)
        # end if

        # Title
        if show_title:
            plt.title(title)
        # end if

        # X labels
        if xlim is not None: plt.xlim(xlim)
        if xticks is not None: plt.xticks(xticks)

        # Y limits
        if ylim is not None: plt.ylim(ylim)
        if yticks is not None: plt.yticks(yticks)
//  [python]
def plot_singular_values(self, point_name, states, idxs, color, linewidth=1, length=-1, show_title=True, title="",
                                   xticks=None, yticks=None, ylim=None, xlim=None, log10=False):
        """
        Plot singular values
        :param point_name: Name of the observation point
        :param states: Observer's state (or None if all)
        :param idxs: Samples indices in the state collection (or None if all)
        :param color: Line color
        :param linewidth: Line width
        :param length: Length of the plot
        :param show_title: Add title to the plot ?
        :param title: Plot title
        :param xticks: X ticks
        :param yticks: Y ticks
        :param ylim: Y limits [min, max]
        :param xlim: X limits [min, max]
        :param log10: If true, show log10 of singular values
        """
        # Get observed data
        cell_states = self._get_observer_data(point_name, states, idxs, concat=True, concat_dim=2)

        # Correlation matrix R
        R = torch.mm(cell_states.t(), cell_states) / cell_states.size(0)

        # SVD on state
        _, S, _ = torch.linalg.svd(R)

        # Log10?
        if log10:
            S = torch.log10(S)
        # end if

        # Learning PC energy
        if length != -1:
            plt.plot(S[:length], color=color, linewidth=linewidth)
        else:
            plt.plot(S, color=color, linewidth=linewidth)
        # end if

        # Title
        if show_title:
            plt.title(title)
        # end if

        # X labels
        if xlim is not None: plt.xlim(xlim)
        if xticks is not None: plt.xticks(xticks)

        # Y limits
        if ylim is not None: plt.ylim(ylim)
        if yticks is not None: plt.yticks(yticks)
//  [python]
def decompose_affine_transformation(module, X_in, y, steps):
    A = module._parameters.get('weight').data  # get tensor out of Parameter
    b = module._parameters.get('bias').data  # get tensor out of Parameter
    U, S, V = torch.svd(A)  # rotation, non-uniform scaling, rotation
    S = torch.diag(S)  # A = U @ S @ V'
    steps = int(steps / 2)
    I = torch.eye(2)

    # Full affine transformation
    # X_out = X_in @ A' + b = X_in @ V @ S @ U' + b

    # Rotation 1
    X_out = interpolate_rot(X_in, y, steps, V)
    X_in = X_out

    # Scaling
    X_in = torch.cat((I * 0, I, X_in))
    X_out = X_in @ S
    interpolate(X_in, X_out, y, steps, 'Scaling', plotting_bases=True)
    X_in = X_out[4:]

    # Rotation 2
    X_out = interpolate_rot(X_in, y, steps, U.t())
    X_in = X_out

    # Translation
    X_in = torch.cat((I * 0, I, X_in))
    X_out = X_in + b
    interpolate(X_in, X_out, y, steps, 'Translating', plotting_bases=True)

    return X_out[4:]
//  [python]
def decompose_affine_transformation(module, X_in, y, steps):
    A = module._parameters.get('weight').data  # get tensor out of Parameter
    b = module._parameters.get('bias').data  # get tensor out of Parameter
    U, S, V = torch.linalg.svd(A)  # rotation, non-uniform scaling, rotation
    S = torch.diag(S)  # A = U @ S @ V'
    steps = int(steps / 2)
    I = torch.eye(2)

    # Full affine transformation
    # X_out = X_in @ A' + b = X_in @ V @ S @ U' + b

    # Rotation 1
    X_out = interpolate_rot(X_in, y, steps, V)
    X_in = X_out

    # Scaling
    X_in = torch.cat((I * 0, I, X_in))
    X_out = X_in @ S
    interpolate(X_in, X_out, y, steps, 'Scaling', plotting_bases=True)
    X_in = X_out[4:]

    # Rotation 2
    X_out = interpolate_rot(X_in, y, steps, U.t())
    X_in = X_out

    # Translation
    X_in = torch.cat((I * 0, I, X_in))
    X_out = X_in + b
    interpolate(X_in, X_out, y, steps, 'Translating', plotting_bases=True)

    return X_out[4:]
//  [python]
def pinv(mat: T.FloatTensor) -> T.FloatTensor:
    """
    Compute matrix pseudoinverse.

    Args:
        mat: A square matrix.

    Returns:
        tensor: The matrix pseudoinverse.

    """
    U, s, V = torch.svd(mat)
    S = unsqueeze(s.reciprocal(), axis=0)
    return multiply(V,S).mm(U.t())
//  [python]
def pinv(mat: T.FloatTensor) -> T.FloatTensor:
    """
    Compute matrix pseudoinverse.

    Args:
        mat: A square matrix.

    Returns:
        tensor: The matrix pseudoinverse.

    """
    U, s, V = torch.linalg.svd(mat)
    S = unsqueeze(s.reciprocal(), axis=0)
    return multiply(V,S).mm(U.t())
//  [python]
def svd(mat: T.FloatTensor) -> T.Tuple[T.FloatTensor]:
    """
    Compute the Singular Value decomposition of a matrix
    A = U S V^T

    Args:
        mat: A matrix.

    Returns:
        (U, S, V): Tuple of tensors.

    """
    return torch.svd(mat)
//  [python]
def svd(mat: T.FloatTensor) -> T.Tuple[T.FloatTensor]:
    """
    Compute the Singular Value decomposition of a matrix
    A = U S V^T

    Args:
        mat: A matrix.

    Returns:
        (U, S, V): Tuple of tensors.

    """
    return torch.linalg.svd(mat)
//  [python]
def fetch_pc(emb, weights, train_loader, pc_path, RESET):
    sentences = torch.cat([d[1][0] for d in train_loader]).int()
    emb_dataset = apply_weights(emb, weights, sentences)

    if not os.path.exists(pc_path) or RESET:
        _, _, V = torch.svd(emb_dataset - emb_dataset.mean(dim=0), some=True)
        v = V[:, 0].unsqueeze(-1)
        u = v.mm(v.t())
        with open(pc_path, 'wb') as file:
            pickle.dump(u, file)
    else:
        with open(pc_path, 'rb') as file:
            u = pickle.load(file)
    return u
//  [python]
def fetch_pc(emb, weights, train_loader, pc_path, RESET):
    sentences = torch.cat([d[1][0] for d in train_loader]).int()
    emb_dataset = apply_weights(emb, weights, sentences)

    if not os.path.exists(pc_path) or RESET:
        _, _, V = torch.linalg.svd(emb_dataset - emb_dataset.mean(dim=0))
        v = V[:, 0].unsqueeze(-1)
        u = v.mm(v.t())
        with open(pc_path, 'wb') as file:
            pickle.dump(u, file)
    else:
        with open(pc_path, 'rb') as file:
            u = pickle.load(file)
    return u
//  [python]
def forward(
            self,
            p: Tensor,
            q: Tensor) -> Tensor:
        ''' Aligns two point clouds using the optimal R, T

            Parameters
            ----------
                p: BxNx3, torch.Tensor
                    The first of points
                q: BxNx3, torch.Tensor

            Returns
            -------
                p_hat: BxNx3, torch.Tensor
                    The points p after least squares alignment to q
        '''
        batch_size = p.shape[0]
        dtype = p.dtype
        device = p.device

        p_transpose = p.transpose(1, 2)
        q_transpose = q.transpose(1, 2)

        # 1. Remove mean.
        p_mean = torch.mean(p_transpose, dim=-1, keepdim=True)
        q_mean = torch.mean(q_transpose, dim=-1, keepdim=True)

        p_centered = p_transpose - p_mean
        q_centered = q_transpose - q_mean

        # 2. Compute variance of X1 used for scale.
        var_p = torch.sum(p_centered.pow(2), dim=(1, 2), keepdim=True)
        #  var_q = torch.sum(q_centered.pow(2), dim=(1, 2), keepdim=True)

        # Compute the outer product of the two point sets
        # Should be Bx3x3
        K = torch.bmm(p_centered, q_centered.transpose(1, 2))
        # Apply SVD on the outer product matrix to recover the rotation
        U, S, V = torch.svd(K)

        # Make sure that the computed rotation does not contain a reflection
        Z = torch.eye(3, dtype=dtype, device=device).view(
            1, 3, 3).expand(batch_size, -1, -1).contiguous()

        raw_product = torch.bmm(U, V.transpose(1, 2))
        Z[:, -1, -1] *= torch.sign(torch.det(raw_product))

        # Compute the final rotation matrix
        rotation = torch.bmm(V, torch.bmm(Z, U.transpose(1, 2)))

        scale = torch.einsum('bii->b', [torch.bmm(rotation, K)]) / var_p.view(
            -1)

        # Compute the translation vector
        translation = q_mean - scale.reshape(batch_size, 1, 1) * torch.bmm(
            rotation, p_mean)

        return (
            scale.reshape(batch_size, 1, 1) *
            torch.bmm(rotation, p_transpose) +
            translation).transpose(1, 2)
//  [python]
def forward(
            self,
            p: Tensor,
            q: Tensor) -> Tensor:
        ''' Aligns two point clouds using the optimal R, T

            Parameters
            ----------
                p: BxNx3, torch.Tensor
                    The first of points
                q: BxNx3, torch.Tensor

            Returns
            -------
                p_hat: BxNx3, torch.Tensor
                    The points p after least squares alignment to q
        '''
        batch_size = p.shape[0]
        dtype = p.dtype
        device = p.device

        p_transpose = p.transpose(1, 2)
        q_transpose = q.transpose(1, 2)

        # 1. Remove mean.
        p_mean = torch.mean(p_transpose, dim=-1, keepdim=True)
        q_mean = torch.mean(q_transpose, dim=-1, keepdim=True)

        p_centered = p_transpose - p_mean
        q_centered = q_transpose - q_mean

        # 2. Compute variance of X1 used for scale.
        var_p = torch.sum(p_centered.pow(2), dim=(1, 2), keepdim=True)
        #  var_q = torch.sum(q_centered.pow(2), dim=(1, 2), keepdim=True)

        # Compute the outer product of the two point sets
        # Should be Bx3x3
        K = torch.bmm(p_centered, q_centered.transpose(1, 2))
        # Apply SVD on the outer product matrix to recover the rotation
        U, S, V = torch.linalg.svd(K)

        # Make sure that the computed rotation does not contain a reflection
        Z = torch.eye(3, dtype=dtype, device=device).view(
            1, 3, 3).expand(batch_size, -1, -1).contiguous()

        raw_product = torch.bmm(U, V.transpose(1, 2))
        Z[:, -1, -1] *= torch.sign(torch.det(raw_product))

        # Compute the final rotation matrix
        rotation = torch.bmm(V, torch.bmm(Z, U.transpose(1, 2)))

        scale = torch.einsum('bii->b', [torch.bmm(rotation, K)]) / var_p.view(
            -1)

        # Compute the translation vector
        translation = q_mean - scale.reshape(batch_size, 1, 1) * torch.bmm(
            rotation, p_mean)

        return (
            scale.reshape(batch_size, 1, 1) *
            torch.bmm(rotation, p_transpose) +
            translation).transpose(1, 2)
//  [python]
def whiten_and_color(self, cF, sF):
        cFSize = cF.size()
        c_mean = torch.mean(cF,1) # c x (h x w)
        c_mean = c_mean.unsqueeze(1).expand_as(cF)
        cF = cF - c_mean
        contentConv = torch.mm(cF,cF.t()).div(cFSize[1]-1) + torch.eye(cFSize[0]).double()
        c_u,c_e,c_v = torch.svd(contentConv,some=False)

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

        sFSize = sF.size()
        s_mean = torch.mean(sF,1)
        sF = sF - s_mean.unsqueeze(1).expand_as(sF)
        styleConv = torch.mm(sF,sF.t()).div(sFSize[1]-1)
        s_u,s_e,s_v = torch.svd(styleConv,some=False)

        k_s = sFSize[0]
        for i in range(sFSize[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)
        targetFeature = torch.mm(torch.mm(torch.mm(s_v[:,0:k_s],torch.diag(s_d)),(s_v[:,0:k_s].t())),whiten_cF)
        targetFeature = targetFeature + s_mean.unsqueeze(1).expand_as(targetFeature)
        return targetFeature
//  [python]
def whiten_and_color(self, cF, sF):
        cFSize = cF.size()
        c_mean = torch.mean(cF,1) # c x (h x w)
        c_mean = c_mean.unsqueeze(1).expand_as(cF)
        cF = cF - c_mean
        contentConv = torch.mm(cF,cF.t()).div(cFSize[1]-1) + torch.eye(cFSize[0]).double()
        c_u,c_e,c_v = torch.linalg.svd(contentConv)

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

        sFSize = sF.size()
        s_mean = torch.mean(sF,1)
        sF = sF - s_mean.unsqueeze(1).expand_as(sF)
        styleConv = torch.mm(sF,sF.t()).div(sFSize[1]-1)
        s_u,s_e,s_v = torch.linalg.svd(styleConv)

        k_s = sFSize[0]
        for i in range(sFSize[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)
        targetFeature = torch.mm(torch.mm(torch.mm(s_v[:,0:k_s],torch.diag(s_d)),(s_v[:,0:k_s].t())),whiten_cF)
        targetFeature = targetFeature + s_mean.unsqueeze(1).expand_as(targetFeature)
        return targetFeature
//  [python]
def test_svd(A):
    u, d, v = geoopt.linalg.svd(A)
    with torch.no_grad():
        for i, a in enumerate(A):
            ut, dt, vt = torch.svd(a)
            np.testing.assert_allclose(u.detach()[i], ut.detach())
            np.testing.assert_allclose(d.detach()[i], dt.detach())
            np.testing.assert_allclose(v.detach()[i], vt.detach())
    u.sum().backward()
//  [python]
def test_svd(A):
    u, d, v = geoopt.linalg.svd(A)
    with torch.no_grad():
        for i, a in enumerate(A):
            ut, dt, vt = torch.linalg.svd(a)
            np.testing.assert_allclose(u.detach()[i], ut.detach())
            np.testing.assert_allclose(d.detach()[i], dt.detach())
            np.testing.assert_allclose(v.detach()[i], vt.detach())
    u.sum().backward()
//  [python]
def U(self):
        '''Rotation matrix'''
        # center
        xyz1 = self.xyz1 - self.centroid1
        xyz2 = self.xyz2 - self.centroid2

        # Computation of the covariance matrix
        C = torch.matmul(xyz1.permute(0,2,1), xyz2)

        # Compute optimal rotation matrix using SVD
        try:
            V, S, W = torch.svd(C.to(torch.float32))
        except: #incase ill conditioned doesnt work if full of nans
            V, S, W = torch.svd(C.to(torch.float32)+1e-2*C.mean()*torch.rand(C.shape, device=C.device))


        # get sign to ensure right-handedness
        d = torch.ones([self.B,3,3], device=xyz1.device)
        d[:,:,-1] = torch.sign(torch.det(V)*torch.det(W)).unsqueeze(1)

        # Rotation matrix U
        U = torch.matmul(d*V, W.permute(0,2,1)) # (B, 3, 3)
        
        if xyz1.dtype == torch.float16: #set ref to half
            U = U.to(torch.float16)

        return U
//  [python]
def U(self):
        '''Rotation matrix'''
        # center
        xyz1 = self.xyz1 - self.centroid1
        xyz2 = self.xyz2 - self.centroid2

        # Computation of the covariance matrix
        C = torch.matmul(xyz1.permute(0,2,1), xyz2)

        # Compute optimal rotation matrix using SVD
        try:
            U, S, V = torch.linalg.svd(C.to(torch.float32))
        except: #incase ill conditioned doesnt work if full of nans
            U, S, V = torch.linalg.svd(C.to(torch.float32)+1e-2*C.mean()*torch.rand(C.shape, device=C.device))


        # get sign to ensure right-handedness
        d = torch.ones([self.B,3,3], device=xyz1.device)
        d[:,:,-1] = torch.sign(torch.det(U)*torch.det(V)).unsqueeze(1)

        # Rotation matrix U
        U = torch.matmul(d*U, V.permute(0,2,1)) # (B, 3, 3)
        
        if xyz1.dtype == torch.float16: #set ref to half
            U = U.to(torch.float16)

        return U
//  [python]
def rotation_matrix(xyz1, xyz2):
    '''
    Inputs
    -----------
    xyz1 (torch.tensor, (B, n_atom, 3))
    '''
    
    B = xyz1.shape[0]

    # center
    xyz1 -= xyz1.mean(1)
    xyz2 -= xyz2.mean(1)

    # Computation of the covariance matrix
    C = torch.matmul(xyz1.permute(0,2,1), xyz2)

    # Compute optimal rotation matrix using SVD
    V, S, W = torch.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([B,3,3], device=xyz1.device)
    d[:,:,-1] = torch.sign(torch.det(V)*torch.det(W)).unsqueeze(1)

    # Rotation matrix U
    U = torch.matmul(d*V, W.permute(0,2,1)) # (B, 3, 3)

    return U
//  [python]
def rotation_matrix(xyz1, xyz2):
    '''
    Inputs
    -----------
    xyz1 (torch.tensor, (B, n_atom, 3))
    '''
    
    B = xyz1.shape[0]

    # center
    xyz1 -= xyz1.mean(1)
    xyz2 -= xyz2.mean(1)

    # Computation of the covariance matrix
    C = torch.matmul(xyz1.permute(0,2,1), xyz2)

    # Compute optimal rotation matrix using SVD
    U, S, V = torch.linalg.svd(C)

    # get sign to ensure right-handedness
    d = torch.ones([B,3,3], device=xyz1.device)
    d[:,:,-1] = torch.sign(torch.det(U)*torch.det(V)).unsqueeze(1)

    # Rotation matrix U
    U = torch.matmul(d*U, V.permute(0,2,1)) # (B, 3, 3)

    return U
//  [python]
def normalize_rotation_matrix(R: torch.Tensor, check_det=True):
    r"""
    Normalize rotation matrices using svd. (torch, batch)

    :param R: Rotation matrix tensor that can reshape to [batch_size, 3, 3].
    :param check_det: Check if the determinant is 1.
    :return: Normalized rotation matrix tensor of shape [batch_size, 3, 3].
    """
    R = R.view(-1, 3, 3)
    U, S, V = torch.svd(R)
    R = U.bmm(V.transpose(1, 2))
    if check_det:
        m = R.det() < 0
        R[m] = U[m].matmul(torch.diag(torch.tensor([1, 1, -1.], device=R.device))).bmm(V[m].transpose(1, 2))
    return R
//  [python]
def normalize_rotation_matrix(R: torch.Tensor, check_det=True):
    r"""
    Normalize rotation matrices using svd. (torch, batch)

    :param R: Rotation matrix tensor that can reshape to [batch_size, 3, 3].
    :param check_det: Check if the determinant is 1.
    :return: Normalized rotation matrix tensor of shape [batch_size, 3, 3].
    """
    R = R.view(-1, 3, 3)
    U, S, V = torch.linalg.svd(R)
    R = U.bmm(V.transpose(1, 2))
    if check_det:
        m = R.det() < 0
        R[m] = U[m].matmul(torch.diag(torch.tensor([1, 1, -1.], device=R.device))).bmm(V[m].transpose(1, 2))
    return R
//  [python]
def __init__(self, dim, num_heads = 4, dim_head = None, discriminator = False):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        self.dim_head = int(dim / num_heads) if dim_head is None else dim_head
        self.weight_dim = self.num_heads * self.dim_head
        self.to_qkv = nn.Linear(dim, self.weight_dim * 3, bias = False)
        self.scale_factor = dim ** -0.5
        self.discriminator = discriminator
        self.w_out = nn.Linear(self.weight_dim, dim, bias = True)

        if discriminator:
            u, s, v = torch.svd(self.to_qkv.weight)
            self.init_spect_norm = torch.max(s)
//  [python]
def __init__(self, dim, num_heads = 4, dim_head = None, discriminator = False):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        self.dim_head = int(dim / num_heads) if dim_head is None else dim_head
        self.weight_dim = self.num_heads * self.dim_head
        self.to_qkv = nn.Linear(dim, self.weight_dim * 3, bias = False)
        self.scale_factor = dim ** -0.5
        self.discriminator = discriminator
        self.w_out = nn.Linear(self.weight_dim, dim, bias = True)

        if discriminator:
            u, s, v = torch.linalg.svd(self.to_qkv.weight)
            self.init_spect_norm = torch.max(s)
//  [python]
def cpu_svd(X):
    # it is much faster to run symeig on CPU
    X = X.cpu() 
    U,D,V_t = torch.svd(X)
    U = U.cuda()
    D = D.cuda()
    V_t = V_t.cuda() 
    return U,D,V_t
//  [python]
def cpu_svd(X):
    # it is much faster to run symeig on CPU
    X = X.cpu() 
    U, S, Vh = torch.linalg.svd(X)
    U = U.cuda()
    S = S.cuda()
    Vh = Vh.cuda() 
    return U, S, Vh
//  [python]
def forward(self, x):
        assert x.dim() == 3

        if self.discriminator:
            u, s, v = torch.svd(self.to_qkv.weight)
            self.to_qkv.weight = torch.nn.Parameter(self.to_qkv.weight * self.init_spect_norm / torch.max(s))

        # Generate the q, k, v vectors
        qkv = self.to_qkv(x)
        q, k, v = tuple(rearrange(qkv, 'b t (d k h) -> k b h t d', k = 3, h = self.num_heads))

        # Enforcing Lipschitzness of Transformer Discriminator
        # Due to Lipschitz constant of standard dot product self-attention
        # layer can be unbounded, so adopt the l2 attention replace the dot product.
        if self.discriminator:
            attn = torch.cdist(q, k, p = 2)
        else:
            attn = torch.einsum("... i d, ... j d -> ... i j", q, k)
        scale_attn = attn * self.scale_factor
        scale_attn_score = torch.softmax(scale_attn, dim = -1)
        result = torch.einsum("... i j, ... j d -> ... i d", scale_attn_score, v)

        # re-compose
        result = rearrange(result, "b h t d -> b t (h d)")
        return self.w_out(result)
//  [python]
def forward(self, x):
        assert x.dim() == 3

        if self.discriminator:
            u, s, v = torch.linalg.svd(self.to_qkv.weight)
            self.to_qkv.weight = torch.nn.Parameter(self.to_qkv.weight * self.init_spect_norm / torch.max(s))

        # Generate the q, k, v vectors
        qkv = self.to_qkv(x)
        q, k, v = tuple(rearrange(qkv, 'b t (d k h) -> k b h t d', k = 3, h = self.num_heads))

        # Enforcing Lipschitzness of Transformer Discriminator
        # Due to Lipschitz constant of standard dot product self-attention
        # layer can be unbounded, so adopt the l2 attention replace the dot product.
        if self.discriminator:
            attn = torch.cdist(q, k, p = 2)
        else:
            attn = torch.einsum("... i d, ... j d -> ... i j", q, k)
        scale_attn = attn * self.scale_factor
        scale_attn_score = torch.softmax(scale_attn, dim = -1)
        result = torch.einsum("... i j, ... j d -> ... i d", scale_attn_score, v)

        # re-compose
        result = rearrange(result, "b h t d -> b t (h d)")
        return self.w_out(result)
//  [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 matSqrt(x):
    U,D,V = torch.svd(x)
    return U * (D.pow(0.5).diag()) * V.t()
//  [python]
def matSqrt(x):
    U, S, V = torch.linalg.svd(x)
    return U @ S.sqrt().diag() @ V.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])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=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))

    # 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 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])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=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))

    # 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 compute_rigid_transform(a: torch.Tensor, b: torch.Tensor, weights: torch.Tensor):
    """Compute rigid transforms between two point sets

    Args:
        a (torch.Tensor): (B, M, 3) points
        b (torch.Tensor): (B, N, 3) points
        weights (torch.Tensor): (B, M)

    Returns:
        Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
    """

    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
    centroid_a = torch.sum(a * weights_normalized, dim=1)
    centroid_b = torch.sum(b * weights_normalized, dim=1)
    a_centered = a - centroid_a[:, None, :]
    b_centered = b - centroid_b[:, None, :]
    cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)

    # Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
    # and choose based on determinant to avoid flips
    u, s, v = torch.svd(cov, some=False, compute_uv=True)
    rot_mat_pos = v @ u.transpose(-1, -2)
    v_neg = v.clone()
    v_neg[:, :, 2] *= -1
    rot_mat_neg = v_neg @ u.transpose(-1, -2)
    rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
    assert torch.all(torch.det(rot_mat) > 0)

    # Compute translation (uncenter centroid)
    translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]

    transform = torch.cat((rot_mat, translation), dim=2)
    return transform
//  [python]
def compute_rigid_transform(a: torch.Tensor, b: torch.Tensor, weights: torch.Tensor):
    """Compute rigid transforms between two point sets

    Args:
        a (torch.Tensor): (B, M, 3) points
        b (torch.Tensor): (B, N, 3) points
        weights (torch.Tensor): (B, M)

    Returns:
        Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
    """

    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
    centroid_a = torch.sum(a * weights_normalized, dim=1)
    centroid_b = torch.sum(b * weights_normalized, dim=1)
    a_centered = a - centroid_a[:, None, :]
    b_centered = b - centroid_b[:, None, :]
    cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)

    # Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
    # and choose based on determinant to avoid flips
    u, s, v = torch.linalg.svd(cov)
    rot_mat_pos = v @ u.transpose(-1, -2)
    v_neg = v.clone()
    v_neg[:, :, 2] *= -1
    rot_mat_neg = v_neg @ u.transpose(-1, -2)
    rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
    assert torch.all(torch.det(rot_mat) > 0)

    # Compute translation (uncenter centroid)
    translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]

    transform = torch.cat((rot_mat, translation), dim=2)
    return transform
//  [python]
def svd(weight):
    """Initialize the layer with the U,V of SVD.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    U,S,V = torch.svd(weight)
    return torch.matmul(U, V.T)
//  [python]
def svd(weight):
    """Initialize the layer with the U,V of SVD.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    U, S, V = torch.linalg.svd(weight)
    return torch.matmul(U, V.T)
//  [python]
def spectral_normalization(weight):
    """Initialize the layer with spectral normalization.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    U,S,V = torch.svd(weight)
    return weight / S.max()
//  [python]
def spectral_normalization(weight):
    """Initialize the layer with spectral normalization.

    Args:
        weight (torch.FloatTensor): Matrix of shape [M, N].

    Returns:
        (torch.FloatTensor): Matrix of shape [M, N].
    """
    U,S,V = torch.linalg.svd(weight)
    return weight / S.max()
//  [python]
def solve_horn_batch(x_1, x_2):
    """Assumes x_1, x_2 are BXNX3, outputs Bx3x3 matrices"""
    

    x_1_n = x_1 - x_1.mean(dim=1, keepdim=True)
    x_2_n = x_2 - x_1.mean(dim=1, keepdim=True)
    
    W = (1./(x_1.shape[0]))*x_2_n.T.dot(x_1_n)

    U,_,V = torch.svd(W)
    S = np.eye(3)
    S[2,2] = np.linalg.det(U) * np.linalg.det(V)
    C = U.dot(S).dot(V)
    return C
//  [python]
def solve_horn_batch(x_1, x_2):
    """Assumes x_1, x_2 are BXNX3, outputs Bx3x3 matrices"""
    

    x_1_n = x_1 - x_1.mean(dim=1, keepdim=True)
    x_2_n = x_2 - x_1.mean(dim=1, keepdim=True)
    
    W = (1./(x_1.shape[0]))*x_2_n.T.dot(x_1_n)

    U,S,V = torch.linalg.svd(W)
    S = np.eye(3)
    S[2,2] = np.linalg.det(U) * np.linalg.det(V)
    C = U.dot(S).dot(V)
    return C
//  [python]
def PCA(data, k=15):
    X = torch.from_numpy(data)
    X_mean = torch.mean(X, 0)
    X = X - X_mean.expand_as(X)
    U, S, V = torch.svd(torch.t(X))
    return torch.mm(X, U[:, :k])
//  [python]
def PCA(data, k=15):
    X = torch.from_numpy(data)
    X_mean = torch.mean(X, 0)
    X = X - X_mean.expand_as(X)
    U, S, V = torch.linalg.svd(torch.t(X))
    return torch.mm(X, U[:, :k])
//  [python]
def svd(mat: torch.Tensor) -> torch.Tensor:
    device = mat.device
    batch_size = mat.size(0)
    dtype = mat.dtype

    if mat.ndim == 2:
        mat = mat.reshape(batch_size, 3, 3)
    if dtype != torch.double: # important for numerical stability
        mat = mat.double()

    u, _, v = torch.svd(mat)

    # Handle reflection via determinant (1 or -1)
    d = torch.det(u @ v.transpose(1, 2))
    d = torch.cat([
        torch.ones(batch_size, 2, device=device),
        d.unsqueeze(-1)
    ], axis=-1)
    d = torch.eye(3, device=device).unsqueeze(0) * d.view(-1, 1, 3)

    # Predict and return a batch of 3x3 orthogonal matrices
    mat = u @ d @ v.transpose(1, 2)
    return mat.to(dtype)
//  [python]
def svd(mat: torch.Tensor) -> torch.Tensor:
    device = mat.device
    batch_size = mat.size(0)
    dtype = mat.dtype

    if mat.ndim == 2:
        mat = mat.reshape(batch_size, 3, 3)
    if dtype != torch.double: # important for numerical stability
        mat = mat.double()

    u, s, v = torch.linalg.svd(mat)

    # Handle reflection via determinant (1 or -1)
    d = torch.det(u @ v.transpose(1, 2))
    d = torch.cat([
        torch.ones(batch_size, 2, device=device),
        d.unsqueeze(-1)
    ], axis=-1)
    d = torch.eye(3, device=device).unsqueeze(0) * d.view(-1, 1, 3)

    # Predict and return a batch of 3x3 orthogonal matrices
    mat = u @ d @ v.transpose(1, 2)
    return mat.to(dtype)
//  [python]
def calc_rmsd(pred, true, mask):
    # pred (N,B,Lasu,natom,3)
    # true (B,Lasu,natom,3)
    # mask (B,Lasu,natom)
    def rmsd(V, W, eps=1e-6):
        L = V.shape[1]
        return torch.sqrt(torch.sum((V-W)*(V-W), dim=(1,2)) / L + eps)
    def centroid(X):
        return X.mean(dim=-2, keepdim=True)

    N, B, L, Natm = pred.shape[:4]
    resmask = mask[0,:,1]
    pred = pred[:,:,resmask,1].squeeze(1)
    true = true[:,resmask,1]
    cP = centroid(pred)
    cT = centroid(true)
    pred = pred - cP
    true = true - cT
    C = torch.einsum('bji,njk->bik', pred, true)
    V, S, W = torch.svd(C)
    d = torch.ones([N,3,3], device=pred.device)
    d[:,:,-1] = torch.sign(torch.det(V)*torch.det(W)).unsqueeze(1)
    U = torch.matmul(d*V, W.permute(0,2,1)) # (IB, 3, 3)
    rpred = torch.matmul(pred, U) # (IB, L*3, 3)
    rms = rmsd(rpred, true).reshape(N)
    return rms, U, cP, cT
//  [python]
def calc_rmsd(pred, true, mask):
    # pred (N,B,Lasu,natom,3)
    # true (B,Lasu,natom,3)
    # mask (B,Lasu,natom)
    def rmsd(V, W, eps=1e-6):
        L = V.shape[1]
        return torch.sqrt(torch.sum((V-W)*(V-W), dim=(1,2)) / L + eps)
    def centroid(X):
        return X.mean(dim=-2, keepdim=True)

    N, B, L, Natm = pred.shape[:4]
    resmask = mask[0,:,1]
    pred = pred[:,:,resmask,1].squeeze(1)
    true = true[:,resmask,1]
    cP = centroid(pred)
    cT = centroid(true)
    pred = pred - cP
    true = true - cT
    C = torch.einsum('bji,njk->bik', pred, true)
    U, S, V = torch.linalg.svd(C)
    d = torch.ones([N,3,3], device=pred.device)
    d[:,:,-1] = torch.sign(torch.det(U)*torch.det(V)).unsqueeze(1)
    W = torch.matmul(d*U, V.permute(0,2,1)) # (IB, 3, 3)
    rpred = torch.matmul(pred, W) # (IB, L*3, 3)
    rms = rmsd(rpred, true).reshape(N)
    return rms, W, cP, cT