jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def find_rigid_alignment_batch(self, A: torch.Tensor, B: torch.Tensor)->torch.Tensor:
        """
        Baased from: https://gist.github.com/bougui505/e392a371f5bab095a3673ea6f4976cc8
        
        See: https://en.wikipedia.org/wiki/Kabsch_algorithm
        2-D or 3-D registration with known correspondences.
        Registration occurs in the zero centered coordinate system, and then
        must be transported back.
        Args:
        - A: Torch tensor of shape (batch_size, N,D) -- Point Cloud to Align (source)
        - B: Torch tensor of shape (batch_size, N,D) -- Reference Point Cloud (target)
        Returns:
        - R: optimal rotation, shape (batch_size, D,D)
        - t: optimal translation, shape (batch_size, D)
        """
        batch_size = A.shape[0]
        a_mean = A.mean(axis=1)
        b_mean = B.mean(axis=1)
        A_c = A - a_mean[:,None,:]
        B_c = B - b_mean[:,None,:]
        # Covariance matrix
        H = torch.bmm(A_c.transpose(1,2),B_c)
        U, S, V = torch.linalg.svd(H)
        # Rotation matrix
        R = torch.bmm(V,U.transpose(1,2))
        # Translation vector
        t = b_mean[:,None,:] - torch.bmm(R, a_mean.unsqueeze(-1)).transpose(1,2)

        t = t.transpose(1,2)
        
        return R, t.squeeze(2)
//  [python]
def __init__(self,
                #  d_input,
                #  d_output,
                 d_model=128,
                 chunk=1,
                 eps=0.01,
                 beta=0.8,
                 gamma=0.01,
                 gated=False,
                 init_std=1,
                 alpha=1,
                 model='LipschitzRNN',
                 solver='euler',
                 l_output=0,
                 l_max=-1,
                 ):
        super().__init__()

        # self.d_input = d_input
        self.d_model = d_model
        # self.chunk = chunk
        self.eps = eps
        self.model = model
        self.solver = solver
        self.gamma = gamma
        self.beta = beta
        self.alpha = alpha


        self.tanh = nn.Tanh()
        self.sigmoid = nn.Sigmoid()

        # self.E = nn.Linear(d_input*self.chunk, d_model)
        # self.D = nn.Linear(d_model, d_output)
        self.register_buffer('I', torch.eye(d_model))

        if self.model == 'simpleRNN':
            self.W = nn.Linear(d_model, d_model, bias=False)
            self.W.weight.data = gaussian_init_(d_model, std=init_std)

        elif self.model == 'resRNN':
            self.W = nn.Linear(d_model, d_model, bias=False)
            self.W.weight.data = gaussian_init_(d_model, std=init_std)

        elif self.model == 'asymRNN':
            self.C = nn.Parameter(gaussian_init_(d_model, std=init_std))

        elif self.model == 'calRNN':
            U, _, V = torch.svd(gaussian_init_(d_model, std=init_std))
            self.C = nn.Parameter(torch.mm(U, V.t()).float())

        elif self.model == 'LipschitzRNN':
            self.C = nn.Parameter(gaussian_init_(d_model, std=init_std))
            self.B = nn.Parameter(gaussian_init_(d_model, std=init_std))

        elif self.model == 'LipschitzRNN_gated':
            self.C = nn.Parameter(gaussian_init_(d_model, std=init_std))
            self.B = nn.Parameter(gaussian_init_(d_model, std=init_std))
            # self.E_gate = nn.Linear(d_input, d_model)

        elif self.model == 'LipschitzRNN_ODE':
            self.func = LipschitzRNN_ODE(d_model, beta, gamma, init_std)

        else:
            print("Unexpected model!")
            raise NotImplementedError
//  [python]
def __init__(self,
                #  d_input,
                #  d_output,
                 d_model=128,
                 chunk=1,
                 eps=0.01,
                 beta=0.8,
                 gamma=0.01,
                 gated=False,
                 init_std=1,
                 alpha=1,
                 model='LipschitzRNN',
                 solver='euler',
                 l_output=0,
                 l_max=-1,
                 ):
        super().__init__()

        # self.d_input = d_input
        self.d_model = d_model
        # self.chunk = chunk
        self.eps = eps
        self.model = model
        self.solver = solver
        self.gamma = gamma
        self.beta = beta
        self.alpha = alpha


        self.tanh = nn.Tanh()
        self.sigmoid = nn.Sigmoid()

        # self.E = nn.Linear(d_input*self.chunk, d_model)
        # self.D = nn.Linear(d_model, d_output)
        self.register_buffer('I', torch.eye(d_model))

        if self.model == 'simpleRNN':
            self.W = nn.Linear(d_model, d_model, bias=False)
            self.W.weight.data = gaussian_init_(d_model, std=init_std)

        elif self.model == 'resRNN':
            self.W = nn.Linear(d_model, d_model, bias=False)
            self.W.weight.data = gaussian_init_(d_model, std=init_std)

        elif self.model == 'asymRNN':
            self.C = nn.Parameter(gaussian_init_(d_model, std=init_std))

        elif self.model == 'calRNN':
            U, S, V = torch.linalg.svd(gaussian_init_(d_model, std=init_std))
            self.C = nn.Parameter(torch.mm(U, V.t()).float())

        elif self.model == 'LipschitzRNN':
            self.C = nn.Parameter(gaussian_init_(d_model, std=init_std))
            self.B = nn.Parameter(gaussian_init_(d_model, std=init_std))

        elif self.model == 'LipschitzRNN_gated':
            self.C = nn.Parameter(gaussian_init_(d_model, std=init_std))
            self.B = nn.Parameter(gaussian_init_(d_model, std=init_std))
            # self.E_gate = nn.Linear(d_input, d_model)

        elif self.model == 'LipschitzRNN_ODE':
            self.func = LipschitzRNN_ODE(d_model, beta, gamma, init_std)

        else:
            print("Unexpected model!")
            raise NotImplementedError
//  [python]
def solve_R(f1, f2):
    """f1 and f2: (b*)m*3
    only work for batch_size=1
    """
    S = torch.matmul(f1.transpose(-1, -2), f2)  # 3*3
    U, sigma, V = torch.svd(S)
    R = torch.matmul(V, U.transpose(-1, -2))
    det = torch.linalg.det(R)

    diag_1 = torch.tensor([1, 1, 0], device=R.device, dtype=R.dtype)
    diag_2 = torch.tensor([0, 0, 1], device=R.device, dtype=R.dtype)
    det_mat = torch.diag(diag_1 + diag_2 * det)

    det_mat = det_mat.unsqueeze(0)

    R = torch.matmul(V, torch.matmul(det_mat, U.transpose(-1, -2)))
    return R
//  [python]
def solve_R(f1, f2):
    """f1 and f2: (b*)m*3
    only work for batch_size=1
    """
    S = torch.matmul(f1.transpose(-1, -2), f2)  # 3*3
    U, S, V = torch.linalg.svd(S)
    R = torch.matmul(V, U.transpose(-1, -2))
    det = torch.linalg.det(R)

    diag_1 = torch.tensor([1, 1, 0], device=R.device, dtype=R.dtype)
    diag_2 = torch.tensor([0, 0, 1], device=R.device, dtype=R.dtype)
    det_mat = torch.diag(diag_1 + diag_2 * det)

    det_mat = det_mat.unsqueeze(0)

    R = torch.matmul(V, torch.matmul(det_mat, U.transpose(-1, -2)))
    return R
//  [python]
def whiten_and_color(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(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 project(self, base):
        try:
            # Compute the projection using the thin SVD decomposition
            U, _, V = torch.svd(base, some=True)
            return U.mm(V.t())
        except RuntimeError:
            # If the svd does not converge, fallback to the (thin) QR decomposition
            x = base
            if base.size(0) < base.size(1):
                x = base.t()
            ret = torch.qr(x, some=True).Q
            if base.size(0) < base.size(1):
                ret = ret.t()
            return ret
//  [python]
def project(self, base):
        try:
            # Compute the projection using the thin SVD decomposition
            U, _, V = torch.linalg.svd(base, full_matrices=False)
            return U.mm(V.t())
        except RuntimeError:
            # If the svd does not converge, fallback to the (thin) QR decomposition
            x = base
            if base.size(0) < base.size(1):
                x = base.t()
            ret = torch.qr(x, some=True).Q
            if base.size(0) < base.size(1):
                ret = ret.t()
            return ret
//  [python]
def pca_points(self, patch_points):
        '''

        Args:
            patch_points: xyz points

        Returns:
            patch_points: xyz points after aligning using pca
        '''
        # compute pca of points in the patch:
        # center the patch around the mean:
        pts_mean = patch_points.mean(0)
        patch_points = patch_points - pts_mean

        trans, _, _ = torch.svd(torch.t(patch_points))
        patch_points = torch.mm(patch_points, trans)

        cp_new = -pts_mean  # since the patch was originally centered, the original cp was at (0,0,0)
        cp_new = torch.matmul(cp_new, trans)

        # re-center on original center point
        patch_points = patch_points - cp_new
        return patch_points, trans
//  [python]
def pca_points(self, patch_points):
        '''

        Args:
            patch_points: xyz points

        Returns:
            patch_points: xyz points after aligning using pca
        '''
        # compute pca of points in the patch:
        # center the patch around the mean:
        pts_mean = patch_points.mean(0)
        patch_points = patch_points - pts_mean

        trans, _, _ = torch.linalg.svd(torch.t(patch_points))
        patch_points = torch.mm(patch_points, trans)

        cp_new = -pts_mean  # since the patch was originally centered, the original cp was at (0,0,0)
        cp_new = torch.matmul(cp_new, trans)

        # re-center on original center point
        patch_points = patch_points - cp_new
        return patch_points, trans
//  [python]
def low_rank(target, sparsity):
    N, target_l, seq_l = target.shape
    target_lr = target.clone()
    try:
        u, s, v = torch.svd(target_lr)
        topk = int(round(seq_l * (1 - sparsity)))
        # assert torch.dist(target_lr, torch.matmul(torch.matmul(u, torch.diag_embed(s)), v.transpose(-2, -1)))<1e-2
        s[:, topk:] = 0
        target_lr = torch.matmul(torch.matmul(u, torch.diag_embed(s)), v.transpose(-2, -1))
        return target_lr, True
    except:  # torch.svd may have convergence issues for GPU and CPU.
        return target_lr, False
//  [python]
def low_rank(target, sparsity):
    N, target_l, seq_l = target.shape
    target_lr = target.clone()
    try:
        u, s, v = torch.linalg.svd(target_lr)
        topk = int(round(seq_l * (1 - sparsity)))
        # assert torch.dist(target_lr, torch.matmul(torch.matmul(u, torch.diag_embed(s)), v.transpose(-2, -1)))<1e-2
        s[:, topk:] = 0
        target_lr = torch.matmul(torch.matmul(u, torch.diag_embed(s)), v.transpose(-2, -1))
        return target_lr, True
    except:  # torch.svd may have convergence issues for GPU and CPU.
        return target_lr, False
//  [python]
def triangulate_point_from_multiple_views_linear_torch(proj_matricies, points, confidences=None):
    """Similar as triangulate_point_from_multiple_views_linear() but for PyTorch.
    For more information see its documentation.
    Args:
        proj_matricies torch tensor of shape (N, 3, 4): sequence of projection matricies (3x4)
        points torch tensor of of shape (N, 2): sequence of points' coordinates
        confidences None or torch tensor of shape (N,): confidences of points [0.0, 1.0].
                                                        If None, all confidences are supposed to be 1.0
    Returns:
        point_3d numpy torch tensor of shape (3,): triangulated point
    """
    assert len(proj_matricies) == len(points)

    n_views = len(proj_matricies)

    if confidences is None:
        confidences = torch.ones(n_views, dtype=torch.float32, device=points.device)

    A = proj_matricies[:, 2:3].expand(n_views, 2, 4) * points.view(n_views, 2, 1)
    A -= proj_matricies[:, :2]
    A *= confidences.view(-1, 1, 1)

    u, s, vh = torch.svd(A.view(-1, 4))

    point_3d_homo = -vh[:, 3]
    point_3d = homogeneous_to_euclidean(point_3d_homo.unsqueeze(0))[0]

    return point_3d
//  [python]
def triangulate_point_from_multiple_views_linear_torch(proj_matricies, points, confidences=None):
    """Similar as triangulate_point_from_multiple_views_linear() but for PyTorch.
    For more information see its documentation.
    Args:
        proj_matricies torch tensor of shape (N, 3, 4): sequence of projection matricies (3x4)
        points torch tensor of of shape (N, 2): sequence of points' coordinates
        confidences None or torch tensor of shape (N,): confidences of points [0.0, 1.0].
                                                        If None, all confidences are supposed to be 1.0
    Returns:
        point_3d numpy torch tensor of shape (3,): triangulated point
    """
    assert len(proj_matricies) == len(points)

    n_views = len(proj_matricies)

    if confidences is None:
        confidences = torch.ones(n_views, dtype=torch.float32, device=points.device)

    A = proj_matricies[:, 2:3].expand(n_views, 2, 4) * points.view(n_views, 2, 1)
    A -= proj_matricies[:, :2]
    A *= confidences.view(-1, 1, 1)

    u, s, vh = torch.linalg.svd(A.view(-1, 4))

    point_3d_homo = -vh[:, 3]
    point_3d = homogeneous_to_euclidean(point_3d_homo.unsqueeze(0))[0]

    return point_3d
//  [python]
def PCA_svd(pts):
    mean = pts.mean(dim=0)
    pts_adjust = pts - mean
    H = torch.mm(pts_adjust.transpose(0, 1), pts_adjust)
    eigenvectors, eigenvalues, eigenvectors_T = torch.svd(H)
    return mean, eigenvectors[:,2]
//  [python]
def PCA_svd(pts):
    mean = pts.mean(dim=0)
    pts_adjust = pts - mean
    H = torch.mm(pts_adjust.transpose(0, 1), pts_adjust)
    eigenvectors, eigenvalues, eigenvectors_T = torch.linalg.svd(H)
    return mean, eigenvectors[:,2]
//  [python]
def init_object_orientation(src_axis, tgt_axis):
        # pseudo inverse
        assert len(src_axis.shape) == 3
        tr = torch.bmm(src_axis.transpose(2, 1), src_axis)
        tr_inv = torch.inverse(tr)
        pseudo = torch.bmm(tr_inv, src_axis.transpose(2, 1))

        rot = torch.bmm(pseudo, tgt_axis)  # + 1e-4 * torch.rand(pseudo.shape[0], 3, 3).to(pseudo.device)
        # return rot
        U, S, V = torch.svd(rot)
        R = torch.bmm(U, V.transpose(2, 1))
        return R
//  [python]
def init_object_orientation(src_axis, tgt_axis):
        # pseudo inverse
        assert len(src_axis.shape) == 3
        tr = torch.bmm(src_axis.transpose(2, 1), src_axis)
        tr_inv = torch.inverse(tr)
        pseudo = torch.bmm(tr_inv, src_axis.transpose(2, 1))

        rot = torch.bmm(pseudo, tgt_axis)  # + 1e-4 * torch.rand(pseudo.shape[0], 3, 3).to(pseudo.device)
        # return rot
        U, S, V = torch.linalg.svd(rot)
        R = torch.bmm(U, V.transpose(2, 1))
        return R
//  [python]
def reset_params(self, init_params=None):
        """Reset the parameters.

        Args:
            init_params (torch.Tensor, optional): Initial parameters.

        Returns:
            None.

        """
        mat = torch.randn((1, 2**self.n_wires, 2**self.n_wires), dtype=C_DTYPE)
        U, Sigma, V = torch.svd(mat)
        self.params.data.copy_(U.matmul(V.permute(0, 2, 1)))
//  [python]
def reset_params(self, init_params=None):
        """Reset the parameters.

        Args:
            init_params (torch.Tensor, optional): Initial parameters.

        Returns:
            None.

        """
        mat = torch.randn((1, 2**self.n_wires, 2**self.n_wires), dtype=C_DTYPE)
        U, S, V = torch.linalg.svd(mat)
        self.params.data.copy_(U.matmul(V.permute(0, 2, 1)))
//  [python]
def WeightedKabsch(self, A, B, W):
    
        assert A.size() == B.size()
    
        batch_size, num_rows, num_cols = A.size()
       
        ## mask to 0/1 weights for motive/static points
        #W=W.type(torch.bool).unsqueeze(2)
        W = W.unsqueeze(2)
        # find mean column wise
        centroid_A = torch.sum(A.transpose(2,1).contiguous()*W, axis=1)
        centroid_B = torch.sum(B.transpose(2,1).contiguous()*W, axis=1)
    
        #print(W.sum(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, _, 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)
        #d = torch.zeros(batch_size).type(torch.int8).cuda()
        # -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 WeightedKabsch(self, A, B, W):
    
        assert A.size() == B.size()
    
        batch_size, num_rows, num_cols = A.size()
       
        ## mask to 0/1 weights for motive/static points
        #W=W.type(torch.bool).unsqueeze(2)
        W = W.unsqueeze(2)
        # find mean column wise
        centroid_A = torch.sum(A.transpose(2,1).contiguous()*W, axis=1)
        centroid_B = torch.sum(B.transpose(2,1).contiguous()*W, axis=1)
    
        #print(W.sum(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)
        #d = torch.zeros(batch_size).type(torch.int8).cuda()
        # -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 least_squares(a, b, rcond=None):
    """
    A PyTorch implementation of NumPy's "linalg.lstsq" function

    Parameters
    ----------
    a : torch.Tensor
        (m,n) "Coefficient" matrix
    b : torch.Tensor
        (m,) or (m,k) "dependent variable" values
    rcond : float
        Cutt-off ratio for small singular values of a. For the
        purposes of rank determination, singular values are treated
        as zero if they are smaller than rcond times the largest
        singular value of a. If `None`, the default will use the
        machine precision times max(m, n)

    Returns
    -------
    x : torch.Tensor
        (n,) or (n,k) least-squares solution
    residuals: torch.Tensor
        (1,) or (k,) or (0,) sums of residuals; squared Euclidean 2-norm for
        each column in b - a*x. If rank(a) < n or m <= n, this is an empty
        array. If b is 1-dimensional, this is a (1,) shape array. Otherwise
        the shape is (k,).
    rank : int
        rank of matrix a
    s : torch.Tensor
        (min(m,n),) singular values of matrix a

    """
    m,n = a.shape
    if rcond is None:
        rcond = max(a.shape)*torch.finfo(a.dtype).eps
    U, s, V = torch.svd(a)
    rank = torch.sum(s > rcond*s[0]).item()
    s_inv = torch.where(s > rcond*s[0], s.reciprocal(), torch.zeros_like(s))
    x = torch.matmul(V, s_inv.unsqueeze(1) * torch.matmul(U.T, b))
    if rank < n or m <= n:
        residuals = torch.tensor([])
    else:
        residuals = torch.sum((a@x - b)**2, 0, keepdim=len(b.shape)==1)

    return x, residuals, rank, s
//  [python]
def least_squares(a, b, rcond=None):
    """
    A PyTorch implementation of NumPy's "linalg.lstsq" function

    Parameters
    ----------
    a : torch.Tensor
        (m,n) "Coefficient" matrix
    b : torch.Tensor
        (m,) or (m,k) "dependent variable" values
    rcond : float
        Cutt-off ratio for small singular values of a. For the
        purposes of rank determination, singular values are treated
        as zero if they are smaller than rcond times the largest
        singular value of a. If `None`, the default will use the
        machine precision times max(m, n)

    Returns
    -------
    x : torch.Tensor
        (n,) or (n,k) least-squares solution
    residuals: torch.Tensor
        (1,) or (k,) or (0,) sums of residuals; squared Euclidean 2-norm for
        each column in b - a*x. If rank(a) < n or m <= n, this is an empty
        array. If b is 1-dimensional, this is a (1,) shape array. Otherwise
        the shape is (k,).
    rank : int
        rank of matrix a
    s : torch.Tensor
        (min(m,n),) singular values of matrix a

    """
    m,n = a.shape
    if rcond is None:
        rcond = max(a.shape)*torch.finfo(a.dtype).eps
    U, s, V = torch.linalg.svd(a)
    rank = torch.sum(s > rcond*s[0]).item()
    s_inv = torch.where(s > rcond*s[0], s.reciprocal(), torch.zeros_like(s))
    x = torch.matmul(V, s_inv.unsqueeze(1) * torch.matmul(U.T, b))
    if rank < n or m <= n:
        residuals = torch.tensor([])
    else:
        residuals = torch.sum((a@x - b)**2, 0, keepdim=len(b.shape)==1)

    return x, residuals, rank, s
//  [python]
def subspaces_principal_angle(vs, us):  # vs, us ~ [batch, n, N]
    vs, us = orthogonalize(vs).transpose(-1, -2), orthogonalize(us).transpose(-1, -2)
    M = vs.transpose(-1, -2) @ us
    svd = torch.svd(M)[1]
    return svd[:, 0] if len(svd.shape) == 2 else svd[0]
//  [python]
def subspaces_principal_angle(vs, us):  # vs, us ~ [batch, n, N]
    vs, us = orthogonalize(vs).transpose(-1, -2), orthogonalize(us).transpose(-1, -2)
    M = vs.transpose(-1, -2) @ us
    svd = torch.linalg.svd(M)[1]
    return svd[:, 0] if len(svd.shape) == 2 else svd[0]
//  [python]
def procrustes_align(
        features_src: torch.Tensor,
        features_tgt: torch.Tensor,
        beta: float = 0.85) -> torch.Tensor:

    u, _, v = torch.svd(features_src.T @ features_tgt)
    W = u @ v.T

    identity = torch.eye(W.size(0)).to(W.device)
    W = W - (W - identity) * beta
    return W
//  [python]
def procrustes_align(
        features_src: torch.Tensor,
        features_tgt: torch.Tensor,
        beta: float = 0.85) -> torch.Tensor:

    u, s, v = torch.linalg.svd(features_src.T @ features_tgt)
    W = u @ v.T

    identity = torch.eye(W.size(0)).to(W.device)
    W = W - (W - identity) * beta
    return W
//  [python]
def spectral_projection(trasfm: torch.Tensor) -> torch.Tensor:
    u, s, v = torch.svd(trasfm)
    s[s > 1] = 1.0
    s[s < 0] = 0.0
    trasfm = u @ (torch.diag(s) @ v.T)
    return trasfm
//  [python]
def spectral_projection(trasfm: torch.Tensor) -> torch.Tensor:
    u, s, v = torch.linalg.svd(trasfm)
    s[s > 1] = 1.0
    s[s < 0] = 0.0
    trasfm = u @ (torch.diag(s) @ v.T)
    return trasfm
//  [python]
def __init__(self,
                 input_len=96,
                 pred_len=96,
                 dynamic_dim=128,
                 encoder=None,
                 decoder=None):
        super(TimeInvKP, self).__init__()
        self.dynamic_dim = dynamic_dim
        self.input_len = input_len
        self.pred_len = pred_len
        self.encoder = encoder
        self.decoder = decoder

        K_init = torch.randn(self.dynamic_dim, self.dynamic_dim)
        U, _, V = torch.svd(K_init) # stable initialization
        self.K = nn.Linear(self.dynamic_dim, self.dynamic_dim, bias=False)
        self.K.weight.data = torch.mm(U, V.t())
//  [python]
def __init__(self,
                 input_len=96,
                 pred_len=96,
                 dynamic_dim=128,
                 encoder=None,
                 decoder=None):
        super(TimeInvKP, self).__init__()
        self.dynamic_dim = dynamic_dim
        self.input_len = input_len
        self.pred_len = pred_len
        self.encoder = encoder
        self.decoder = decoder

        K_init = torch.randn(self.dynamic_dim, self.dynamic_dim)
        U, S, V = torch.linalg.svd(K_init) # stable initialization
        self.K = nn.Linear(self.dynamic_dim, self.dynamic_dim, bias=False)
        self.K.weight.data = torch.mm(U, V.t())
//  [python]
def forward(self, f_s, f_t):
        _, s_s, _ = torch.svd(f_s)
        _, s_t, _ = torch.svd(f_t)
        loss = torch.pow(s_s[0], 2) + torch.pow(s_t[0], 2)
        return loss
//  [python]
def forward(self, f_s, f_t):
        _, s_s, _ = torch.linalg.svd(f_s)
        _, s_t, _ = torch.linalg.svd(f_t)
        loss = torch.pow(s_s[0], 2) + torch.pow(s_t[0], 2)
        return loss
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        logSs = torch.zeros_like(input)
        invSs = torch.zeros_like(input)
        for i in range(input.shape[0]):
            U, S, V = torch.svd(input[i, :, :])
            Ss[i, :] = S
            Us[i, :, :] = U
            logSs[i, :, :] = torch.diag(torch.log(S))
            invSs[i, :, :] = torch.diag(1.0 / S)
        result = torch.matmul(Us, torch.matmul(logSs, torch.transpose(Us, 1, 2)))
        self.Us = Us
        self.Ss = Ss
        self.logSs = logSs
        self.invSs = invSs
        self.save_for_backward(input)
        return result
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        logSs = torch.zeros_like(input)
        invSs = torch.zeros_like(input)
        for i in range(input.shape[0]):
            U, S, V = torch.linalg.svd(input[i, :, :])
            Ss[i, :] = S
            Us[i, :, :] = U
            logSs[i, :, :] = torch.diag(torch.log(S))
            invSs[i, :, :] = torch.diag(1.0 / S)
        result = torch.matmul(Us, torch.matmul(logSs, torch.transpose(Us, 1, 2)))
        self.Us = Us
        self.Ss = Ss
        self.logSs = logSs
        self.invSs = invSs
        self.save_for_backward(input)
        return result
//  [python]
def test_svd_complex(self, device, dtype):
        # this test verifies that torch.svd really returns V and not V.conj()
        # see: https://github.com/pytorch/pytorch/issues/45821
        t = torch.randn((10, 10), dtype=dtype, device=device)
        U, S, V = torch.svd(t, some=False)
        # verify that t ≈ t2
        # t2 = U @ diag(S) @ Vᴴ
        # Vᴴ is the conjugate transpose of V
        t2 = U @ torch.diag(S).type(dtype) @ V.conj().T
        self.assertEqual(t, t2)
//  [python]
def test_svd_complex(self, device, dtype):
        # this test verifies that torch.linalg.svd really returns V and not V.conj()
        # see: https://github.com/pytorch/pytorch/issues/45821
        t = torch.randn((10, 10), dtype=dtype, device=device)
        U, S, V = torch.linalg.svd(t, full_matrices=False)
        # verify that t ≈ t2
        # t2 = U @ diag(S) @ Vᴴ
        # Vᴴ is the conjugate transpose of V
        t2 = U @ torch.diag(S).type(dtype) @ V.conj().T
        self.assertEqual(t, t2)
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        max_Ss = torch.zeros_like(input)
        max_Ids = torch.zeros_like(input)
        for i in range(input.shape[0]):
            U, S, V = torch.svd(input[i, :, :])
            eps = 0.0001
            max_S = torch.clamp(S, min=eps)
            max_Id = torch.ge(S, eps).float()
            Ss[i, :] = S
            Us[i, :, :] = U
            max_Ss[i, :, :] = torch.diag(max_S)
            max_Ids[i, :, :] = torch.diag(max_Id)

        result = torch.matmul(Us, torch.matmul(max_Ss, torch.transpose(Us, 1, 2)))
        self.Us = Us
        self.Ss = Ss
        self.max_Ss = max_Ss
        self.max_Ids = max_Ids
        self.save_for_backward(input)
        return result
//  [python]
def forward(self, input):
        Us = torch.zeros_like(input)
        Ss = torch.zeros((input.shape[0], input.shape[1])).double()
        max_Ss = torch.zeros_like(input)
        max_Ids = torch.zeros_like(input)
        for i in range(input.shape[0]):
            U, S, V = torch.linalg.svd(input[i, :, :])
            eps = 0.0001
            max_S = torch.clamp(S, min=eps)
            max_Id = torch.ge(S, eps).float()
            Ss[i, :] = S
            Us[i, :, :] = U
            max_Ss[i, :, :] = torch.diag(max_S)
            max_Ids[i, :, :] = torch.diag(max_Id)

        result = torch.matmul(Us, torch.matmul(max_Ss, torch.transpose(Us, 1, 2)))
        self.Us = Us
        self.Ss = Ss
        self.max_Ss = max_Ss
        self.max_Ids = max_Ids
        self.save_for_backward(input)
        return result
//  [python]
def iteration_step(self, batch, instance_params, weight_dict={}):
        """
        Computes losses for a single step of optimization.
        Entries in loss/weight dict should have the following entries (always edit loss_keys to modify loss terms):
        corr, templ, s2m, m2s, pose_pr, shape_pr
        """
        loss_keys = ['corr', 'templ', 's2m_h', 's2m_o', 'm2s', 'pose_pr', 'shape_pr', 'smpl_J', 'pca_axis']
        for k in loss_keys:
            if k not in weight_dict.keys():
                weight_dict[k] = 1.

        device = self.device
        loss = {}
        inputs = batch.get('inputs').to(device)
        p = batch.get('grid_coords').to(device)
        points = batch.get('points').to(device)
        joints = batch.get('smpl_J').to(device)
        pc_h = batch.get('pc_h').to(device)
        pc_o = batch.get('pc_o').to(device)
        batch_sz = inputs.shape[0]

        poses, betas, trans = instance_params['pose'], instance_params['betas'], instance_params['trans']
        corr, df_h = instance_params['corr'], instance_params['df_h']
        # import ipdb; ipdb.set_trace()
        ## FIT OBJECT
        object_init = instance_params['object_init']
        rot = instance_params['obj_R']

        U, S, V = torch.svd(rot + 1e-3 * torch.rand(batch_sz, 3, 3).to(device))
        R = torch.bmm(U, V.transpose(2, 1))
        object = torch.bmm(object_init, R) + instance_params['obj_t'].unsqueeze(1)

        # evaluate object DF at object points
        out_ = self.model(object, inputs)
        loss['object'] = out_['df_o'].mean() * weight_dict['object']

        # S2M on human and object
        loss['s2m_o'], _ = chamfer_distance(pc_o, object)
        loss['s2m_o'] *= weight_dict['s2m_o']

        posed_scan_correspondences = self.vsmpl(corr, poses, betas, trans)

        # SMPL joints
        _, smpl_J, _, _ = self.smpl_layer(poses, betas, trans)
        # loss on 3D joints lifted from the images.
        loss['smpl_J'] = F.l1_loss(smpl_J, joints) * weight_dict['smpl_J']

        # correspondence loss
        loss['corr'] = F.l1_loss(torch.linalg.norm(points - posed_scan_correspondences, dim=-1), df_h,
                                 reduction='none').mean() * weight_dict['corr']

        # pose prior
        loss['pose_pr'] = self.pose_prior(poses).mean() * weight_dict['pose_pr']

        # shape prior
        loss['shape_pr'] = (betas ** 2).mean() * weight_dict['shape_pr']

        return loss
//  [python]
def iteration_step(self, batch, instance_params, weight_dict={}):
        """
        Computes losses for a single step of optimization.
        Entries in loss/weight dict should have the following entries (always edit loss_keys to modify loss terms):
        corr, templ, s2m, m2s, pose_pr, shape_pr
        """
        loss_keys = ['corr', 'templ', 's2m_h', 's2m_o', 'm2s', 'pose_pr', 'shape_pr', 'smpl_J', 'pca_axis']
        for k in loss_keys:
            if k not in weight_dict.keys():
                weight_dict[k] = 1.

        device = self.device
        loss = {}
        inputs = batch.get('inputs').to(device)
        p = batch.get('grid_coords').to(device)
        points = batch.get('points').to(device)
        joints = batch.get('smpl_J').to(device)
        pc_h = batch.get('pc_h').to(device)
        pc_o = batch.get('pc_o').to(device)
        batch_sz = inputs.shape[0]

        poses, betas, trans = instance_params['pose'], instance_params['betas'], instance_params['trans']
        corr, df_h = instance_params['corr'], instance_params['df_h']
        # import ipdb; ipdb.set_trace()
        ## FIT OBJECT
        object_init = instance_params['object_init']
        rot = instance_params['obj_R']

        U, S, V = torch.linalg.svd(rot + 1e-3 * torch.rand(batch_sz, 3, 3).to(device))
        R = torch.bmm(U, V.transpose(2, 1))
        object = torch.bmm(object_init, R) + instance_params['obj_t'].unsqueeze(1)

        # evaluate object DF at object points
        out_ = self.model(object, inputs)
        loss['object'] = out_['df_o'].mean() * weight_dict['object']

        # S2M on human and object
        loss['s2m_o'], _ = chamfer_distance(pc_o, object)
        loss['s2m_o'] *= weight_dict['s2m_o']

        posed_scan_correspondences = self.vsmpl(corr, poses, betas, trans)

        # SMPL joints
        _, smpl_J, _, _ = self.smpl_layer(poses, betas, trans)
        # loss on 3D joints lifted from the images.
        loss['smpl_J'] = F.l1_loss(smpl_J, joints) * weight_dict['smpl_J']

        # correspondence loss
        loss['corr'] = F.l1_loss(torch.linalg.norm(points - posed_scan_correspondences, dim=-1), df_h,
                                 reduction='none').mean() * weight_dict['corr']

        # pose prior
        loss['pose_pr'] = self.pose_prior(poses).mean() * weight_dict['pose_pr']

        # shape prior
        loss['shape_pr'] = (betas ** 2).mean() * weight_dict['shape_pr']

        return loss
//  [python]
def orthogonalise(mat):
    """Orthogonalise rotation/affine matrices

    Ideally, 3D rotation matrices should be orthogonal,
    however during creation, floating point errors can build up.
    We SVD decompose our matrix as in the ideal case S is a diagonal matrix of 1s
    We then round the values of S to [-1, 0, +1],
    making U @ S_rounded @ V.T an orthonormal matrix close to the original.
    """
    orth_mat = mat.clone()
    u, s, v = torch.svd(mat[..., :3, :3])
    orth_mat[..., :3, :3] = u @ torch.diag_embed(s.round()) @ v.transpose(-1, -2)
    return orth_mat
//  [python]
def orthogonalise(mat):
    """Orthogonalise rotation/affine matrices

    Ideally, 3D rotation matrices should be orthogonal,
    however during creation, floating point errors can build up.
    We SVD decompose our matrix as in the ideal case S is a diagonal matrix of 1s
    We then round the values of S to [-1, 0, +1],
    making U @ S_rounded @ V.T an orthonormal matrix close to the original.
    """
    orth_mat = mat.clone()
    u, s, v = torch.linalg.svd(mat[..., :3, :3])
    orth_mat[..., :3, :3] = u @ torch.diag_embed(s.round()) @ v.transpose(-1, -2)
    return orth_mat
//  [python]
def __init__(self, kernel, channels, img_dim, device, ZERO = 3e-2):
        self.img_dim = img_dim
        self.channels = channels
        #build 1D conv matrix
        H_small = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                H_small[i, j] = kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.svd(H_small, some=False)
        #ZERO = 3e-2
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1), self.singulars_small.reshape(1, img_dim)).reshape(img_dim**2)
        #sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True)
//  [python]
def __init__(self, kernel, channels, img_dim, device, ZERO = 3e-2):
        self.img_dim = img_dim
        self.channels = channels
        #build 1D conv matrix
        H_small = torch.zeros(img_dim, img_dim, device=device)
        for i in range(img_dim):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                if j < 0 or j >= img_dim: continue
                H_small[i, j] = kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.linalg.svd(H_small)
        #ZERO = 3e-2
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1), self.singulars_small.reshape(1, img_dim)).reshape(img_dim**2)
        #sort the big matrix singulars and save the permutation
        self._singulars, self._perm = self._singulars.sort(descending=True)
//  [python]
def test_svd(self):
        a = torch.Tensor(((8.79, 6.11, -9.15, 9.57, -3.49, 9.84),
                          (9.93, 6.91, -7.93, 1.64, 4.02, 0.15),
                          (9.83, 5.04, 4.86, 8.83, 9.80, -8.99),
                          (5.45, -0.27, 4.85, 0.74, 10.00, -6.02),
                          (3.16, 7.98, 3.01, 5.80, 4.27, -5.31))).t().clone()
        u, s, v = torch.svd(a)
        uu = torch.Tensor()
        ss = torch.Tensor()
        vv = torch.Tensor()
        uuu, sss, vvv = torch.svd(a, out=(uu, ss, vv))
        self.assertEqual(u, uu, 0, 'torch.svd')
        self.assertEqual(u, uuu, 0, 'torch.svd')
        self.assertEqual(s, ss, 0, 'torch.svd')
        self.assertEqual(s, sss, 0, 'torch.svd')
        self.assertEqual(v, vv, 0, 'torch.svd')
        self.assertEqual(v, vvv, 0, 'torch.svd')

        # test reuse
        X = torch.randn(4, 4)
        U, S, V = torch.svd(X)
        Xhat = torch.mm(U, torch.mm(S.diag(), V.t()))
        self.assertEqual(X, Xhat, 1e-8, 'USV\' wrong')

        self.assertFalse(U.is_contiguous(), 'U is contiguous')
        torch.svd(X, out=(U, S, V))
        Xhat = torch.mm(U, torch.mm(S.diag(), V.t()))
        self.assertEqual(X, Xhat, 1e-8, 'USV\' wrong')

        # test non-contiguous
        X = torch.randn(5, 5)
        U = torch.zeros(5, 2, 5)[:, 1]
        S = torch.zeros(5, 2)[:, 1]
        V = torch.zeros(5, 2, 5)[:, 1]

        self.assertFalse(U.is_contiguous(), 'U is contiguous')
        self.assertFalse(S.is_contiguous(), 'S is contiguous')
        self.assertFalse(V.is_contiguous(), 'V is contiguous')
        torch.svd(X, out=(U, S, V))
        Xhat = torch.mm(U, torch.mm(S.diag(), V.t()))
        self.assertEqual(X, Xhat, 1e-8, 'USV\' wrong')
//  [python]
def test_svd(self):
        a = torch.Tensor(((8.79, 6.11, -9.15, 9.57, -3.49, 9.84),
                          (9.93, 6.91, -7.93, 1.64, 4.02, 0.15),
                          (9.83, 5.04, 4.86, 8.83, 9.80, -8.99),
                          (5.45, -0.27, 4.85, 0.74, 10.00, -6.02),
                          (3.16, 7.98, 3.01, 5.80, 4.27, -5.31))).t().clone()
        u, s, v = torch.linalg.svd(a)
        uu = torch.Tensor()
        ss = torch.Tensor()
        vv = torch.Tensor()
        uuu, sss, vvv = torch.linalg.svd(a, out=(uu, ss, vv))
        self.assertEqual(u, uu, 0, 'torch.linalg.svd')
        self.assertEqual(u, uuu, 0, 'torch.linalg.svd')
        self.assertEqual(s, ss, 0, 'torch.linalg.svd')
        self.assertEqual(s, sss, 0, 'torch.linalg.svd')
        self.assertEqual(v, vv, 0, 'torch.linalg.svd')
        self.assertEqual(v, vvv, 0, 'torch.linalg.svd')

        # test reuse
        X = torch.randn(4, 4)
        U, S, V = torch.linalg.svd(X)
        Xhat = torch.mm(U, torch.mm(S.diag(), V.t()))
        self.assertEqual(X, Xhat, 1e-8, 'USV\' wrong')

        self.assertFalse(U.is_contiguous(), 'U is contiguous')
        torch.linalg.svd(X, out=(U, S, V))
        Xhat = torch.mm(U, torch.mm(S.diag(), V.t()))
        self.assertEqual(X, Xhat, 1e-8, 'USV\' wrong')

        # test non-contiguous
        X = torch.randn(5, 5)
        U = torch.zeros(5, 2, 5)[:, 1]
        S = torch.zeros(5, 2)[:, 1]
        V = torch.zeros(5, 2, 5)[:, 1]

        self.assertFalse(U.is_contiguous(), 'U is contiguous')
        self.assertFalse(S.is_contiguous(), 'S is contiguous')
        self.assertFalse(V.is_contiguous(), 'V is contiguous')
        torch.linalg.svd(X, out=(U, S, V))
        Xhat = torch.mm(U, torch.mm(S.diag(), V.t()))
        self.assertEqual(X, Xhat, 1e-8, 'USV\' wrong')
//  [python]
def forward(ctx, input):
        # Note: input is matrix of size m x n with m >= n.
        # Note: if above assumption is voilated, the gradients
        # will be wrong.
        U, S, V = torch.svd(input, some=True)

        ctx.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def forward(ctx, input):
        # Note: input is matrix of size m x n with m >= n.
        # Note: if above assumption is voilated, the gradients
        # will be wrong.
        U, S, V = torch.linalg.svd(input, some=True)

        ctx.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def __call__(self, features):
        from sklearn.decomposition import TruncatedSVD
        from scipy.stats import entropy
        import torch

        if isinstance(features, torch.Tensor):
            _,s,_ = torch.svd(features)
            s     = s.cpu().numpy()
        else:
            svd = TruncatedSVD(n_components=self.embed_dim-1, n_iter=7, random_state=42)
            svd.fit(features)
            s = svd.singular_values_

        if self.mode!=0:
            s = s[np.abs(self.mode)-1:]
        s_norm  = s/np.sum(s)
        uniform = np.ones(len(s))/(len(s))

        if self.mode<0:
            kl = entropy(s_norm, uniform)
        if self.mode>0:
            kl = entropy(uniform, s_norm)
        if self.mode==0:
            kl = s_norm

        return kl
//  [python]
def __call__(self, features):
        from sklearn.decomposition import TruncatedSVD
        from scipy.stats import entropy
        import torch

        if isinstance(features, torch.Tensor):
            U, S, Vh = torch.linalg.svd(features)
            s     = S.cpu().numpy()
        else:
            svd = TruncatedSVD(n_components=self.embed_dim-1, n_iter=7, random_state=42)
            svd.fit(features)
            s = svd.singular_values_

        if self.mode!=0:
            s = s[np.abs(self.mode)-1:]
        s_norm  = s/np.sum(s)
        uniform = np.ones(len(s))/(len(s))

        if self.mode<0:
            kl = entropy(s_norm, uniform)
        if self.mode>0:
            kl = entropy(uniform, s_norm)
        if self.mode==0:
            kl = s_norm

        return kl
//  [python]
def solve(self, b, regul=1e-8):
        u, s, v = torch.svd(self.data.view(-1, self.data.size(2)))
        x = torch.mv(v, torch.mv(v.t(), b.get_flat_representation()) /
                     (s**2 + regul))
        return PVector(b.layer_collection, vector_repr=x)
//  [python]
def solve(self, b, regul=1e-8):
        u, s, v = torch.linalg.svd(self.data.view(-1, self.data.size(2)))
        x = torch.mv(v, torch.mv(v.t(), b.get_flat_representation()) /
                     (s**2 + regul))
        return PVector(b.layer_collection, vector_repr=x)
//  [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 batch_svd(A):
    """Wrapper around torch.svd that works when the input is a batch of matrices."""
    U_list = []
    S_list = []
    V_list = []
    for i in range(A.shape[0]):
        U, S, V = torch.svd(A[i])
        U_list.append(U)
        S_list.append(S)
        V_list.append(V)
    U = torch.stack(U_list, dim=0)
    S = torch.stack(S_list, dim=0)
    V = torch.stack(V_list, dim=0)
    return U, S, V
//  [python]
def batch_svd(A):
    """Wrapper around torch.svd that works when the input is a batch of matrices."""
    U_list = []
    S_list = []
    V_list = []
    for i in range(A.shape[0]):
        U, S, V = torch.linalg.svd(A[i])
        U_list.append(U)
        S_list.append(S)
        V_list.append(V)
    U = torch.stack(U_list, dim=0)
    S = torch.stack(S_list, dim=0)
    V = torch.stack(V_list, dim=0)
    return U, S, V
//  [python]
def get_lowrank_tuple_torch_cpu_old(tensor, max_rank):
    u, s, v = torch.svd(tensor.float())
    u, s, v = u[:,:max_rank], s[:max_rank], v.t()[:max_rank, :]
    l       = torch.matmul(u, torch.diag(s))
    del u, s
    return  (v.t(), l.t())
//  [python]
def get_lowrank_tuple_torch_cpu_old(tensor, max_rank):
    u, s, v = torch.linalg.svd(tensor.float())
    u, s, v = u[:,:max_rank], s[:max_rank], v.t()[:max_rank, :]
    l       = torch.matmul(u, torch.diag(s))
    del u, s
    return  (v.t(), l.t())
//  [python]
def mx_inv(mx):
    U, D, V = torch.svd(mx)
    eps = 0.009
    D_min = torch.min(D)
    if D_min < eps:
        D_1 = torch.zeros_like(D)
        D_1[D>D_min] = 1/D[D>D_min]
    else:
        D_1 = 1/D
    #D_1 = 1 / D #.clamp(min=0.005)
    
    return U @ D_1.diag() @ V.t()
//  [python]
def mx_inv(mx):
    U, S, V = torch.linalg.svd(mx)
    eps = 0.009
    D_min = torch.min(S)
    if D_min < eps:
        D_1 = torch.zeros_like(S)
        D_1[S>D_min] = 1/S[S>D_min]
    else:
        D_1 = 1/S
    #D_1 = 1 / D #.clamp(min=0.005)
    
    return U @ D_1.diag() @ V.t()