jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
gpu_timer = Timer()
trials = 10
num_corr = 400

Covariances = torch.abs(torch.rand([num_corr, 3, 3]))
for i in range(trials):
    cpu_timer.tic()
    u, s, vt = torch.svd(Covariances)
    cpu_timer.toc()
//  [python]
gpu_timer = Timer()
trials = 10
num_corr = 400

Covariances = torch.abs(torch.rand([num_corr, 3, 3]))
for i in range(trials):
    cpu_timer.tic()
    u, s, vt = torch.linalg.svd(Covariances)
    cpu_timer.toc()
//  [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 SVD(self, layer, percent_removed):

        W = layer.weight.data
        U, S, V = torch.svd(W)
        L = int((1.-percent_removed)*U.shape[0])
        W1 = U[:,:L]
        W2 = torch.diag(S[:L]) @ V[:,:L].t()
        layer_1 = nn.Linear(in_features=layer.in_features, 
                    out_features=L, bias=False)
        layer_1.weight.data = W2

        layer_2 = nn.Linear(in_features=L, 
                    out_features=layer.out_features, bias=True)
        layer_2.weight.data = W1

        if layer.bias.data is None: 
            layer_2.bias.data = torch.zeros(*layer.out_features.shape)
        else:
            layer_2.bias.data = layer.bias.data

        return nn.Sequential(layer_1, layer_2)
//  [python]
def SVD(self, layer, percent_removed):

        W = layer.weight.data
        U, S, V = torch.linalg.svd(W)
        L = int((1.-percent_removed)*U.shape[0])
        W1 = U[:,:L]
        W2 = torch.diag(S[:L]) @ V[:,:L].t()
        layer_1 = nn.Linear(in_features=layer.in_features, 
                    out_features=L, bias=False)
        layer_1.weight.data = W2

        layer_2 = nn.Linear(in_features=L, 
                    out_features=layer.out_features, bias=True)
        layer_2.weight.data = W1

        if layer.bias.data is None: 
            layer_2.bias.data = torch.zeros(*layer.out_features.shape)
        else:
            layer_2.bias.data = layer.bias.data

        return nn.Sequential(layer_1, layer_2)
//  [python]
def cluster_loss(H, kloss, lmbda, batch_size):
    gram_matrix = (H.T @ H) / batch_size
    _ ,sv_2, _ = torch.svd(gram_matrix)
    sv = torch.sqrt(sv_2[:kloss])
    loss = torch.sum(sv)
    return lmbda*loss
//  [python]
def cluster_loss(H, kloss, lmbda, batch_size):
    gram_matrix = (H.T @ H) / batch_size
    _ ,sv_2, _ = torch.linalg.svd(gram_matrix)
    sv = torch.sqrt(sv_2[:kloss])
    loss = torch.sum(sv)
    return lmbda*loss
//  [python]
def _inverse(self, name, M, code_class, code_pos):
        """
        Matrix inverse
        :param name: Name associated with M
        :param M: Matrix to inverse
        :return: Inverse matrix
        """
        if self._debug == Node.DEBUG_TEST or self._debug == Node.DEBUG_OUTPUT:
            # SVD of matrix
            _, S, _ = torch.svd(M)

            # Condition number
            condition_number = torch.log10(S[0] / S[-1])

            # Show condition number
            if self._debug == Node.DEBUG_OUTPUT:
                print(
                    "DEBUG - INFO : Condition number while inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )

            # Bad condition number
            if condition_number > 14:
                print(
                    "DEBUG - WARNING : High condition number while inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )
            # end if
        # end if
        return torch.inverse(M)
//  [python]
def _inverse(self, name, M, code_class, code_pos):
        """
        Matrix inverse
        :param name: Name associated with M
        :param M: Matrix to inverse
        :return: Inverse matrix
        """
        if self._debug == Node.DEBUG_TEST or self._debug == Node.DEBUG_OUTPUT:
            # SVD of matrix
            _, S, _ = torch.linalg.svd(M)

            # Condition number
            condition_number = torch.log10(S[0] / S[-1])

            # Show condition number
            if self._debug == Node.DEBUG_OUTPUT:
                print(
                    "DEBUG - INFO : Condition number while inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )

            # Bad condition number
            if condition_number > 14:
                print(
                    "DEBUG - WARNING : High condition number while inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )
            # end if
        # end if
        return torch.inverse(M)
//  [python]
def _pinverse(self, name, M, code_class, code_pos):
        """
        Matrix pseudo-inverse
        :param name: Name associated with M
        :param M: Matrix to inverse
        :return: Pseudo-inverse of matrix
        """
        if self._debug == Node.DEBUG_TEST or self._debug == Node.DEBUG_OUTPUT:
            # SVD of matrix
            _, S, _ = torch.svd(M)

            # Condition number
            condition_number = torch.log10(S[0] / S[-1])

            # Show condition number
            if self._debug == Node.DEBUG_OUTPUT:
                print(
                    "DEBUG - INFO : Condition number while pseudo-inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )

            # Bad condition number
            if condition_number > 14:
                print(
                    "DEBUG - WARNING : High condition number while pseudo-inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )
            # end if
        # end if
        return torch.pinverse(M)
//  [python]
def _pinverse(self, name, M, code_class, code_pos):
        """
        Matrix pseudo-inverse
        :param name: Name associated with M
        :param M: Matrix to inverse
        :return: Pseudo-inverse of matrix
        """
        if self._debug == Node.DEBUG_TEST or self._debug == Node.DEBUG_OUTPUT:
            # SVD of matrix
            _, S, _ = torch.linalg.svd(M)

            # Condition number
            condition_number = torch.log10(S[0] / S[-1])

            # Show condition number
            if self._debug == Node.DEBUG_OUTPUT:
                print(
                    "DEBUG - INFO : Condition number while pseudo-inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )

            # Bad condition number
            if condition_number > 14:
                print(
                    "DEBUG - WARNING : High condition number while pseudo-inversing {} : {} (at {}:{})".format(
                        name,
                        condition_number,
                        code_class,
                        code_pos
                    )
                )
            # end if
        # end if
        return torch.pinverse(M)
//  [python]
def get_camera_matrix(self, camera):
        # phimatrix = camera[0:9].view(3,3)
        # U, S,V = torch.svd(phimatrix)
        # R = torch.mm(U, V.permute(1,0))
        # T = camera[9:,0]
        # scale = (phimatrix/R).mean()
        # T = torch.cat([T, self.offset_z + T[None,0]*0])
        # return scale, R, T

        phimatrix = camera[0:6].view(2, 3)
        U, S, V = torch.svd(phimatrix)
        newS = torch.eye(U.size(1), V.size(1)).type(U.type())
        R = torch.mm(torch.mm(U,newS), V.permute(1,0))
        scale = (phimatrix / R).mean()
        R = torch.cat([R, torch.cross(R[0,:], R[1,:]).view(1,-1)], dim=0)
        T = camera[6:].view(-1)
        return scale, R, T
//  [python]
def get_camera_matrix(self, camera):
        phimatrix = camera[0:6].view(2, 3)
        U, S, V = torch.linalg.svd(phimatrix)
        newS = torch.eye(U.size(1), V.size(1)).type(U.type())
        R = torch.mm(torch.mm(U,newS), V.permute(1,0))
        scale = (phimatrix / R).mean()
        R = torch.cat([R, torch.cross(R[0,:], R[1,:]).view(1,-1)], dim=0)
        T = camera[6:].view(-1)
        return scale, R, T
//  [python]
def svd_feature(features, d=200):
    """Get 200-dimensional node features, to avoid curse of dimensionality"""
    if features.shape[1] <= d:
        return features
    U, S, VT = torch.svd(features)
    res = torch.mm(U[:, 0:d], torch.diag(S[0:d]))
    return res
//  [python]
def svd_feature(features, d=200):
    """Get 200-dimensional node features, to avoid curse of dimensionality"""
    if features.shape[1] <= d:
        return features
    U, S, VT = torch.linalg.svd(features)
    res = torch.mm(U[:, 0:d], torch.diag(S[0:d]))
    return res
//  [python]
def orthogonalize(w):
    u, s, v = torch.svd(w)
    return torch.mm(u, v.t())
//  [python]
def orthogonalize(w):
    u, s, v = torch.linalg.svd(w)
    return torch.mm(u, v.t())
//  [python]
def weighted_icp(src, tgt, weights, _EPS = 1e-8):
    """Compute rigid transforms between two point sets

    Args:
        src (torch.Tensor): (B, M, 3) points
        tgt (torch.Tensor): (B, M, 3) points
        weights (torch.Tensor): (B, M)

    Returns:
        R, t, transformed_src: (B, 3, 3), (B, 3), (B, M, 3)

    Modified from open source code:
        https://github.com/yewzijian/RPMNet/blob/master/src/models/rpmnet.py
    """
    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
    centroid_src = torch.sum(src * weights_normalized, dim=1)
    centroid_tgt = torch.sum(tgt * weights_normalized, dim=1)
    src_centered = src - centroid_src[:, None, :]
    tgt_centered = tgt - centroid_tgt[:, None, :]
    cov = src_centered.transpose(-2, -1) @ (tgt_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_src[:, :, None] + centroid_tgt[:, :, None]
    translation = torch.squeeze(translation, -1)
    transformed_src = batch_transform(src, rot_mat, translation)
    return rot_mat, translation, transformed_src
//  [python]
def weighted_icp(src, tgt, weights, _EPS = 1e-8):
    """Compute rigid transforms between two point sets

    Args:
        src (torch.Tensor): (B, M, 3) points
        tgt (torch.Tensor): (B, M, 3) points
        weights (torch.Tensor): (B, M)

    Returns:
        R, t, transformed_src: (B, 3, 3), (B, 3), (B, M, 3)

    Modified from open source code:
        https://github.com/yewzijian/RPMNet/blob/master/src/models/rpmnet.py
    """
    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
    centroid_src = torch.sum(src * weights_normalized, dim=1)
    centroid_tgt = torch.sum(tgt * weights_normalized, dim=1)
    src_centered = src - centroid_src[:, None, :]
    tgt_centered = tgt - centroid_tgt[:, None, :]
    cov = src_centered.transpose(-2, -1) @ (tgt_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_src[:, :, None] + centroid_tgt[:, :, None]
    translation = torch.squeeze(translation, -1)
    transformed_src = batch_transform(src, rot_mat, translation)
    return rot_mat, translation, transformed_src
//  [python]
def svd_for_factorization(self, w: torch.Tensor):
        assert w.ndim == 2
        u, s, v = torch.svd(w)
        a = torch.matmul(u, torch.diag(torch.sqrt(s)))
        b = torch.matmul(torch.diag(torch.sqrt(s)), v.transpose(0, 1))
        print(a.max(), b.max(), w.max())
        return a, b
//  [python]
def svd_for_factorization(self, w: torch.Tensor):
        assert w.ndim == 2
        u, s, v = torch.linalg.svd(w)
        a = torch.matmul(u, torch.diag(torch.sqrt(s)))
        b = torch.matmul(torch.diag(torch.sqrt(s)), v.transpose(0, 1))
        print(a.max(), b.max(), w.max())
        return a, b
//  [python]
def original_implementation(x, y):
    _, s_s, _ = torch.svd(x)
    _, s_t, _ = torch.svd(y)
    sigma = torch.pow(s_s[0], 2) + torch.pow(s_t[0], 2)
    return sigma
//  [python]
def original_implementation(x, y):
    _, s_s, _ = torch.linalg.svd(x)
    _, s_t, _ = torch.linalg.svd(y)
    sigma = torch.pow(s_s[0], 2) + torch.pow(s_t[0], 2)
    return sigma
//  [python]
def svd(feat, n=1):
    size = feat.shape
    assert len(size) == 4

    x = feat.view(size[0], size[1] * size[2], size[3]).float()
    u, s, v = torch.svd(x)

    u = removenan(u)
    s = removenan(s)
    v = 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(feat, n=1):
    size = feat.shape
    assert len(size) == 4

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

    u = removenan(u)
    s = removenan(s)
    v = 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 solve_camera(self, points3d, points2d, ransac_iteration=10, ransac_samples=100):

        assert len(points3d.size()) == 2, 'does not support batching'
        errors = []
        cams = []
        points3d = points3d.detach()
        points2d = points2d.detach()
        N = points3d.size(0)
        for i in range(ransac_iteration):
            try:
                sample_indices = self.sample_random_indices(ransac_samples, N)
                if len(torch.unique(sample_indices)) < 4:
                    continue
                sample_points3d = points3d[sample_indices, :]
                sample_points2d = points2d[sample_indices, :]
                A, b = self.computeAB(sample_points2d, sample_points3d)
                U, S, V = torch.svd(A, some=True)
                S = S*(S > 1E-3).float() + 1E-4*(S < 1E-3).float()
                psuedoSinv = S.pow(-1) * (S > 1E-3).float()
                Sinv = torch.diag(psuedoSinv)
                # Sinv = torch.diag((S + 1E-4).pow(-1))
                Ainv = torch.mm(V, torch.mm(Sinv, U.permute(1, 0)))
                camera_params = torch.mm(Ainv, b)
                scale, R, T = self.get_camera_matrix(camera_params)
                if torch.isnan(R).sum().item() > 0:
                    continue;

                error = self.compute_err(scale, R, T, points3d, points2d)
                cams.append((scale, R, T))
                errors.append(error)
                
        
            except RuntimeError as e: ## Sometimes the matrices are ill-conditioned. We ignore such examples
                continue

        if len(errors) < ransac_iteration/2:
            pdb.set_trace()
            return None, None, None ## The inputs are ill-conditioned

        assert len(errors) > 0,' You should have atleast one entry per image. Not all the samples should be ill-conditioned'
        errors = torch.stack(errors)  # ransac_iteration
        # cams = torch.stack(cams)  # ransac_iterations x 8

        _, ind = torch.min(errors, dim=0)
        cam_best = cams[ind]  # N x 8
        # Find the cam with the least error, and return it.
        return cam_best[0], cam_best[1], cam_best[2]
//  [python]
def solve_camera(self, points3d, points2d, ransac_iteration=10, ransac_samples=100):

        assert len(points3d.size()) == 2, 'does not support batching'
        errors = []
        cams = []
        points3d = points3d.detach()
        points2d = points2d.detach()
        N = points3d.size(0)
        for i in range(ransac_iteration):
            try:
                sample_indices = self.sample_random_indices(ransac_samples, N)
                if len(torch.unique(sample_indices)) < 4:
                    continue
                sample_points3d = points3d[sample_indices, :]
                sample_points2d = points2d[sample_indices, :]
                A, b = self.computeAB(sample_points2d, sample_points3d)
                U, S, V = torch.linalg.svd(A)
                S = S*(S > 1E-3).float() + 1E-4*(S < 1E-3).float()
                psuedoSinv = S.pow(-1) * (S > 1E-3).float()
                Sinv = torch.diag(psuedoSinv)
                # Sinv = torch.diag((S + 1E-4).pow(-1))
                Ainv = torch.mm(V, torch.mm(Sinv, U.permute(1, 0)))
                camera_params = torch.mm(Ainv, b)
                scale, R, T = self.get_camera_matrix(camera_params)
                if torch.isnan(R).sum().item() > 0:
                    continue;

                error = self.compute_err(scale, R, T, points3d, points2d)
                cams.append((scale, R, T))
                errors.append(error)
                
        
            except RuntimeError as e: ## Sometimes the matrices are ill-conditioned. We ignore such examples
                continue

        if len(errors) < ransac_iteration/2:
            pdb.set_trace()
            return None, None, None ## The inputs are ill-conditioned

        assert len(errors) > 0,' You should have atleast one entry per image. Not all the samples should be ill-conditioned'
        errors = torch.stack(errors)  # ransac_iteration
        # cams = torch.stack(cams)  # ransac_iterations x 8

        _, ind = torch.min(errors, dim=0)
        cam_best = cams[ind]  # N x 8
        # Find the cam with the least error, and return it.
        return cam_best[0], cam_best[1], cam_best[2]
//  [python]
def __init__(self, d_in, d_out, rank):
        super().__init__()
        layer = nn.Linear(d_in, d_out)
        w = layer.weight.data
        self.bias = layer.bias
        u, s, v = torch.svd(w)
        s = torch.diag(s[:rank].sqrt())
        u = u[:, :rank]
        v = v.t()[:rank]
        self.u = nn.Parameter((u @ s).t())
        self.v = nn.Parameter((s @ v).t())
//  [python]
def __init__(self, d_in, d_out, rank):
        super().__init__()
        layer = nn.Linear(d_in, d_out)
        w = layer.weight.data
        self.bias = layer.bias
        u, s, v = torch.linalg.svd(w)
        s = torch.diag(s[:rank].sqrt())
        u = u[:, :rank]
        v = v.t()[:rank]
        self.u = nn.Parameter((u @ s).t())
        self.v = nn.Parameter((s @ v).t())
//  [python]
def kabsch_torch(X, Y, cpu=True):
    """ Kabsch alignment of X into Y. 
        Assumes X,Y are both (Dims x N_points). See below for wrapper.
    """
    device = X.device
    #  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()).detach()
    if cpu: 
        C = C.cpu()
    # Optimal rotation matrix via SVD
    if int(torch.__version__.split(".")[1]) < 8:
        # warning! int torch 1.<8 : W must be transposed
        V, S, W = torch.svd(C)
        W = W.t()
    else: 
        V, S, W = torch.linalg.svd(C)
    
    # 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).to(device)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    return X_, Y_
//  [python]
def kabsch_torch(X, Y, cpu=True):
    """ Kabsch alignment of X into Y. 
        Assumes X,Y are both (Dims x N_points). See below for wrapper.
    """
    device = X.device
    #  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()).detach()
    if cpu: 
        C = C.cpu()
    # Optimal rotation matrix via SVD
    V, S, W = torch.linalg.svd(C)
    
    # 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).to(device)
    # calculate rotations
    X_ = torch.matmul(X_.t(), U).t()
    # return centered and aligned
    return X_, Y_
//  [python]
def process_key_eigen(key, checkpoint, method_type, backbone_weight_dict=None):
    lora_a = checkpoint[key]
    lora_b = checkpoint[key.replace("layer1", "layer2")]
    lora_weight = torch.matmul(lora_a, lora_b).float()
    routing_key = key.replace("expert_lora.layer1", "router.expert_embeddings")
    U_lora, S_lora, V_lora = torch.svd(lora_weight)
    if method_type == "eigen_weighted_vector":
        S_weights = S_lora / torch.sum(S_lora)
    elif method_type == "top1_eigen_weighted_vector":
        S_weights = torch.zeros_like(S_lora)
        S_weights[0] = 1
    elif method_type == "backbone_eigen_weighted_vector":
        pass
    new_value = torch.sum(torch.matmul(U_lora, torch.diag(S_weights)), dim=1)
    return routing_key, new_value
//  [python]
def process_key_eigen(key, checkpoint, method_type, backbone_weight_dict=None):
    lora_a = checkpoint[key]
    lora_b = checkpoint[key.replace("layer1", "layer2")]
    lora_weight = torch.matmul(lora_a, lora_b).float()
    routing_key = key.replace("expert_lora.layer1", "router.expert_embeddings")
    U_lora, S_lora, V_lora = torch.linalg.svd(lora_weight)
    if method_type == "eigen_weighted_vector":
        S_weights = S_lora / torch.sum(S_lora)
    elif method_type == "top1_eigen_weighted_vector":
        S_weights = torch.zeros_like(S_lora)
        S_weights[0] = 1
    elif method_type == "backbone_eigen_weighted_vector":
        pass
    new_value = torch.sum(torch.matmul(U_lora, torch.diag(S_weights)), dim=1)
    return routing_key, new_value
//  [python]
def rt_from_fundamental_mat(self, fmat, K, depth_match):
        # F: [b, 3, 3] K: [b, 3, 3] depth_match: [b ,4, n]
        verify_match = self.rand_sample(depth_match, 200) # [b,4,100]
        K_inv = torch.inverse(K)
        b = fmat.shape[0]
        fmat_ = K.transpose(1,2).bmm(fmat)
        essential_mat = fmat_.bmm(K)
        essential_mat_cpu = essential_mat.cpu()
        U, S, V = torch.svd(essential_mat_cpu)
        U, S, V = U.to(K.get_device()), S.to(K.get_device()), V.to(K.get_device())
        W = torch.from_numpy(np.array([[[0., -1., 0.],[1., 0., 0.],[0., 0., 1.]]])).float().repeat(b,1,1).to(K.get_device())
        # R = UWV^t or UW^tV^t t = U[:,2] the third column of U
        R1 = U.bmm(W).bmm(V.transpose(1,2)) # Do we need matrix determinant sign?
        R1 = torch.sign(torch.det(R1)).unsqueeze(-1).unsqueeze(-1) * R1
        R2 = U.bmm(W.transpose(1,2)).bmm(V.transpose(1,2))
        R2 = torch.sign(torch.det(R2)).unsqueeze(-1).unsqueeze(-1) * R2
        t1 = U[:,:,2].unsqueeze(-1) # The third column
        t2 = -U[:,:,2].unsqueeze(-1) # Inverse direction
        
        iden = torch.cat([torch.eye(3), torch.zeros([3,1])], -1).unsqueeze(0).repeat(b,1,1).to(K.get_device()) # [b,3,4]
        P1 = K.bmm(iden)
        P2_1 = K.bmm(torch.cat([R1, t1], -1))
        P2_2 = K.bmm(torch.cat([R2, t1], -1))
        P2_3 = K.bmm(torch.cat([R1, t2], -1))
        P2_4 = K.bmm(torch.cat([R2, t2], -1))
        P2_c = [P2_1, P2_2, P2_3, P2_4]
        flags = []
        for i in range(4):
            with torch.no_grad():
                inlier_num = self.verifyRT(verify_match, K_inv, P1, P2_c[i])
                flags.append(inlier_num)
        P2_c = torch.stack(P2_c, dim=1) # [B, 4, 3, 4]
        flags = torch.stack(flags, dim=1) # [B, 4]
        idx = torch.argmax(flags, dim=-1, keepdim=True) # [b,1]
        P2 = torch.gather(P2_c, index=idx.unsqueeze(-1).unsqueeze(-1).repeat(1,1,3,4), dim=1).squeeze(1) # [b,3,4]
        #pdb.set_trace()
        return P1, P2
//  [python]
def rt_from_fundamental_mat(self, fmat, K, depth_match):
        # F: [b, 3, 3] K: [b, 3, 3] depth_match: [b ,4, n]
        verify_match = self.rand_sample(depth_match, 200) # [b,4,100]
        K_inv = torch.inverse(K)
        b = fmat.shape[0]
        fmat_ = K.transpose(1,2).bmm(fmat)
        essential_mat = fmat_.bmm(K)
        essential_mat_cpu = essential_mat.cpu()
        U, S, V = torch.linalg.svd(essential_mat_cpu)
        U, S, V = U.to(K.get_device()), S.to(K.get_device()), V.to(K.get_device())
        W = torch.from_numpy(np.array([[[0., -1., 0.],[1., 0., 0.],[0., 0., 1.]]])).float().repeat(b,1,1).to(K.get_device())
        # R = UWV^t or UW^tV^t t = U[:,2] the third column of U
        R1 = U.bmm(W).bmm(V.transpose(1,2)) # Do we need matrix determinant sign?
        R1 = torch.sign(torch.det(R1)).unsqueeze(-1).unsqueeze(-1) * R1
        R2 = U.bmm(W.transpose(1,2)).bmm(V.transpose(1,2))
        R2 = torch.sign(torch.det(R2)).unsqueeze(-1).unsqueeze(-1) * R2
        t1 = U[:,:,2].unsqueeze(-1) # The third column
        t2 = -U[:,:,2].unsqueeze(-1) # Inverse direction
        
        iden = torch.cat([torch.eye(3), torch.zeros([3,1])], -1).unsqueeze(0).repeat(b,1,1).to(K.get_device()) # [b,3,4]
        P1 = K.bmm(iden)
        P2_1 = K.bmm(torch.cat([R1, t1], -1))
        P2_2 = K.bmm(torch.cat([R2, t1], -1))
        P2_3 = K.bmm(torch.cat([R1, t2], -1))
        P2_4 = K.bmm(torch.cat([R2, t2], -1))
        P2_c = [P2_1, P2_2, P2_3, P2_4]
        flags = []
        for i in range(4):
            with torch.no_grad():
                inlier_num = self.verifyRT(verify_match, K_inv, P1, P2_c[i])
                flags.append(inlier_num)
        P2_c = torch.stack(P2_c, dim=1) # [B, 4, 3, 4]
        flags = torch.stack(flags, dim=1) # [B, 4]
        idx = torch.argmax(flags, dim=-1, keepdim=True) # [b,1]
        P2 = torch.gather(P2_c, index=idx.unsqueeze(-1).unsqueeze(-1).repeat(1,1,3,4), dim=1).squeeze(1) # [b,3,4]
        #pdb.set_trace()
        return P1, P2
//  [python]
def manual_svd_translation(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    # """Compute the translation vector that aligns A to B using SVD."""
    assert A.size(1) == B.size(1)
    u, s, vt = torch.svd((B.T @ A).T)
    R = u @ vt.T
    return R, s
//  [python]
def manual_svd_translation(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    # """Compute the translation vector that aligns A to B using SVD."""
    assert A.size(1) == B.size(1)
    u, s, vt = torch.linalg.svd((B.T @ A).T)
    R = u @ vt.T
    return R, s
//  [python]
def __init__(self, d_in, d_out, rank=None):
        super().__init__()
        if rank is None:
            self.conv = nn.Conv1d(d_in, d_out, 1)
            self.factorized = False
        else:
            layer = nn.Linear(d_in, d_out)
            w = layer.weight.data
            self.bias = layer.bias
            u, s, v = torch.svd(w)
            s = torch.diag(s[:rank].sqrt())
            u = u[:, :rank]
            v = v.t()[:rank]
            self.u = nn.Parameter(u @ s)
            self.v = nn.Parameter(s @ v)
            self.factorized = True
//  [python]
def __init__(self, d_in, d_out, rank=None):
        super().__init__()
        if rank is None:
            self.conv = nn.Conv1d(d_in, d_out, 1)
            self.factorized = False
        else:
            layer = nn.Linear(d_in, d_out)
            w = layer.weight.data
            self.bias = layer.bias
            u, s, v = torch.linalg.svd(w)
            s = torch.diag(s[:rank].sqrt())
            u = u[:, :rank]
            v = v.t()[:rank]
            self.u = nn.Parameter(u @ s)
            self.v = nn.Parameter(s @ v)
            self.factorized = True
//  [python]
def calc_eigens(args, state_dict):
    # get all of the style modulation weights except for weights in `to_rgb`
    modulated = {
        k: v
        for k, v in state_dict.items()
        if 'style_modulation' in k and 'to_rgb' not in k and 'weight' in k
    }

    weight_mat = []
    for _, v in modulated.items():
        weight_mat.append(v)

    W = torch.cat(weight_mat, dim=0)
    eigen_vector = torch.svd(W).V

    # save eigen vector
    output_path = os.path.splitext(args.ckpt)[0] + '_eigen-vec-mod.pth'
    torch.save({'ckpt': args.ckpt, 'eigen_vector': eigen_vector}, output_path)

    return eigen_vector
//  [python]
def calc_eigens(args, state_dict):
    # get all of the style modulation weights except for weights in `to_rgb`
    modulated = {
        k: v
        for k, v in state_dict.items()
        if 'style_modulation' in k and 'to_rgb' not in k and 'weight' in k
    }

    weight_mat = []
    for _, v in modulated.items():
        weight_mat.append(v)

    W = torch.cat(weight_mat, dim=0)
    eigen_vector = torch.linalg.svd(W).V

    # save eigen vector
    output_path = os.path.splitext(args.ckpt)[0] + '_eigen-vec-mod.pth'
    torch.save({'ckpt': args.ckpt, 'eigen_vector': eigen_vector}, output_path)

    return eigen_vector
//  [python]
def batch_get_pelvis_orient_svd(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    # type: (Tensor, Tensor, Tensor, Tensor, int) -> Tensor
    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]    # this will cause illegal memory access
    ind = torch.where(mask_zero != 0)[0]
    nind = torch.where(mask_zero == 0)[0]
    S_non_zero = S[ind]

    L = S_non_zero.clone()
    U, _, V = torch.svd(L)

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

    det_u_v = torch.sign(det3x3(torch.bmm(V.detach(), U.detach().transpose(1, 2))))
    # det_u_v[det_u_v == 0.0] = 1.0
    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[ind] = 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):
    # type: (Tensor, Tensor, Tensor, Tensor, int) -> Tensor
    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]    # this will cause illegal memory access
    ind = torch.where(mask_zero != 0)[0]
    nind = torch.where(mask_zero == 0)[0]
    S_non_zero = S[ind]

    L = S_non_zero.clone()
    U, S, V = torch.linalg.svd(L)

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

    det_u_v = torch.sign(det3x3(torch.bmm(V, U.transpose(1, 2))))
    # det_u_v[det_u_v == 0.0] = 1.0
    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[ind] = 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):
    # type: (Tensor, Tensor, Tensor, List[int], int) -> Tensor
    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[:, c].clone()
            template = rel_rest_pose[:, c].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))

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

    ind = torch.where(mask_zero != 0)[0]
    nind = torch.where(mask_zero == 0)[0]
    S_non_zero = S[ind]

    L = S_non_zero.clone()
    U, _, V = torch.svd(L)

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

    det_u_v = torch.sign(det3x3(torch.bmm(V.detach(), U.detach().transpose(1, 2))))
    # det_u_v[det_u_v == 0.0] = 1.0
    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[ind] = rot_mat_non_zero

    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):
    # type: (Tensor, Tensor, Tensor, List[int], int) -> Tensor
    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[:, c].clone()
            template = rel_rest_pose[:, c].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))

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

    ind = torch.where(mask_zero != 0)[0]
    nind = torch.where(mask_zero == 0)[0]
    S_non_zero = S[ind]

    L = S_non_zero.clone()
    U, S, V = torch.linalg.svd(L)

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

    det_u_v = torch.sign(det3x3(torch.bmm(V, U.transpose(1, 2))))
    # det_u_v[det_u_v == 0.0] = 1.0
    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[ind] = rot_mat_non_zero

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

    return rot_mat
//  [python]
def prox_nuclear_cuda(self, data, alpha):

        device = data.device
        U, S, V = torch.svd(data)
        # self.nuclear_norm = S.sum()
        # print(f"rank = {len(S.nonzero())}")
        self.nuclear_norm = S.sum()
        S = torch.clamp(S-alpha, min=0)
        indices = torch.tensor([range(0, U.shape[0]),range(0, U.shape[0])]).to(device)
        values = S
        diag_S = torch.sparse.FloatTensor(indices, values, torch.Size(U.shape))
        # diag_S = torch.diag(torch.clamp(S-alpha, min=0))
        # print(f"rank_after = {len(diag_S.nonzero())}")
        V = torch.spmm(diag_S, V.t_())
        V = torch.matmul(U, V)
        return V
//  [python]
def prox_nuclear_cuda(self, data, alpha):

        device = data.device
        U, S, V = torch.linalg.svd(data)
        # self.nuclear_norm = S.sum()
        # print(f"rank = {len(S.nonzero())}")
        self.nuclear_norm = S.sum()
        S = torch.clamp(S-alpha, min=0)
        indices = torch.tensor([range(0, U.shape[0]),range(0, U.shape[0])]).to(device)
        values = S
        diag_S = torch.sparse.FloatTensor(indices, values, torch.Size(U.shape))
        # diag_S = torch.diag(torch.clamp(S-alpha, min=0))
        # print(f"rank_after = {len(diag_S.nonzero())}")
        V = torch.spmm(diag_S, V.t_())
        V = torch.matmul(U, V)
        return V
//  [python]
def whiten(cF):
    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

    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)
    return whiten_cF
//  [python]
def whiten(cF):
    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

    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)
    return whiten_cF
//  [python]
def forward(self, *input):
        src_embedding = input[0]
        tgt_embedding = input[1]
        src = input[2]
        tgt = input[3]
        batch_size = src.size(0)

        d_k = src_embedding.size(1)
        scores = torch.matmul(src_embedding.transpose(2, 1).contiguous(), tgt_embedding) / math.sqrt(d_k)
        scores = torch.softmax(scores, dim=2)

        src_corr = torch.matmul(tgt, scores.transpose(2, 1).contiguous())

        src_centered = src - src.mean(dim=2, keepdim=True)

        src_corr_centered = src_corr - src_corr.mean(dim=2, keepdim=True)

        H = torch.matmul(src_centered, src_corr_centered.transpose(2, 1).contiguous())

        U, S, V = [], [], []
        R = []

        for i in range(src.size(0)):
            u, s, v = torch.svd(H[i])
            r = torch.matmul(v, u.transpose(1, 0).contiguous())
            r_det = torch.det(r)
            if r_det < 0:
                u, s, v = torch.svd(H[i])
                v = torch.matmul(v, self.reflect)
                r = torch.matmul(v, u.transpose(1, 0).contiguous())
                # r = r * self.reflect
            R.append(r)

            U.append(u)
            S.append(s)
            V.append(v)

        U = torch.stack(U, dim=0)
        V = torch.stack(V, dim=0)
        S = torch.stack(S, dim=0)
        R = torch.stack(R, dim=0)

        t = torch.matmul(-R, src.mean(dim=2, keepdim=True)) + src_corr.mean(dim=2, keepdim=True)
        return R, t.view(batch_size, 3)
//  [python]
def forward(self, *input):
        src_embedding = input[0]
        tgt_embedding = input[1]
        src = input[2]
        tgt = input[3]
        batch_size = src.size(0)

        d_k = src_embedding.size(1)
        scores = torch.matmul(src_embedding.transpose(2, 1).contiguous(), tgt_embedding) / math.sqrt(d_k)
        scores = torch.softmax(scores, dim=2)

        src_corr = torch.matmul(tgt, scores.transpose(2, 1).contiguous())

        src_centered = src - src.mean(dim=2, keepdim=True)

        src_corr_centered = src_corr - src_corr.mean(dim=2, keepdim=True)

        H = torch.matmul(src_centered, src_corr_centered.transpose(2, 1).contiguous())

        U, S, V = [], [], []
        R = []

        for i in range(src.size(0)):
            U_i, S_i, V_i = torch.linalg.svd(H[i])
            r = torch.matmul(V_i, U_i.transpose(1, 0).contiguous())
            r_det = torch.det(r)
            if r_det < 0:
                U_i, S_i, V_i = torch.linalg.svd(H[i])
                V_i = torch.matmul(V_i, self.reflect)
                r = torch.matmul(V_i, U_i.transpose(1, 0).contiguous())
                # r = r * self.reflect
            R.append(r)

            U.append(U_i)
            S.append(S_i)
            V.append(V_i)

        U = torch.stack(U, dim=0)
        V = torch.stack(V, dim=0)
        S = torch.stack(S, dim=0)
        R = torch.stack(R, dim=0)

        t = torch.matmul(-R, src.mean(dim=2, keepdim=True)) + src_corr.mean(dim=2, keepdim=True)
        return R, t.view(batch_size, 3)
//  [python]
def compute_rigid_transform(points1, points2, weights):
    """Compute rigid transforms between two point clouds via weighted SVD.
       Adapted from https://github.com/yewzijian/RPMNet/
    Args:
        points1 (torch.Tensor): (B, M, 3) coordinates of the first point cloud
        points2 (torch.Tensor): (B, N, 3) coordinates of the second point cloud
        weights (torch.Tensor): (B, M)
    Returns:
        Transform T (B, 3, 4) to get from points1 to points2, i.e. T*points1 = points2
    """

    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + 1e-5)
    centroid_a = torch.sum(points1 * weights_normalized, dim=1)
    centroid_b = torch.sum(points2 * weights_normalized, dim=1)
    a_centered = points1 - centroid_a[:, None, :]
    b_centered = points2 - 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(points1, points2, weights):
    """Compute rigid transforms between two point clouds via weighted SVD.
       Adapted from https://github.com/yewzijian/RPMNet/
    Args:
        points1 (torch.Tensor): (B, M, 3) coordinates of the first point cloud
        points2 (torch.Tensor): (B, N, 3) coordinates of the second point cloud
        weights (torch.Tensor): (B, M)
    Returns:
        Transform T (B, 3, 4) to get from points1 to points2, i.e. T*points1 = points2
    """

    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + 1e-5)
    centroid_a = torch.sum(points1 * weights_normalized, dim=1)
    centroid_b = torch.sum(points2 * weights_normalized, dim=1)
    a_centered = points1 - centroid_a[:, None, :]
    b_centered = points2 - 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