jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def init_projection_matrix(self, x):
        # Set if using projection matrix
        self.params.use_projection_matrix = getattr(self.params, 'use_projection_matrix', True)

        if self.params.use_projection_matrix:
            self.compressed_dim = self.fparams.attribute('compressed_dim', None)

            proj_init_method = getattr(self.params, 'proj_init_method', 'pca')
            if proj_init_method == 'pca':
                x_mat = TensorList([e.permute(1, 0, 2, 3).reshape(e.shape[1], -1).clone() for e in x])
                x_mat -= x_mat.mean(dim=1, keepdim=True)
                cov_x = x_mat @ x_mat.t()
                self.projection_matrix = TensorList(
                    [None if cdim is None else torch.linalg.svd(C)[0][:, :cdim].t().unsqueeze(-1).unsqueeze(-1).clone() for C, cdim in
                     zip(cov_x, self.compressed_dim)])
            elif proj_init_method == 'randn':
                self.projection_matrix = TensorList(
                    [None if cdim is None else ex.new_zeros(cdim,ex.shape[1],1,1).normal_(0,1/math.sqrt(ex.shape[1])) for ex, cdim in
                     zip(x, self.compressed_dim)])
        else:
            self.compressed_dim = x.size(1)
            self.projection_matrix = TensorList([None]*len(x))
//  [python]
def __init__(self, X, ndim=128, var_fraction=0.99, l2_normalized=True, first_direction=None):

        self.l2_normalized = l2_normalized
        if l2_normalized:
            X = X[:, :-1]

        assert len(X.shape) == 2
        torch.cuda.synchronize()
        start_time = time.time()

        self.mean = torch.mean(X, dim=0, keepdim=True)
        #self.mean = 0
        #self.std = torch.std(X, dim=0, keepdim=True) + 1e-6
        self.std = 1
        #print("std is ", self.std[:, :10].cpu().numpy())
        #X_orig = X
        X = (X - self.mean) / self.std

        U, S, V = torch.svd(X)
        S = S[:ndim]
        V = V[:, :ndim]
        self.proj = V
        scale = torch.mm(X, self.proj).std(dim=0)
        torch.cuda.synchronize()
        print("PCA time taken on vectors of size %s : %f" % (str(X.size()), time.time() - start_time))
        print("largest std of each PC: ", scale[:10].cpu().numpy())
        print("smallest std of each PC: ", scale[-10:].cpu().numpy())
        self.sinvals = S
        print("largest sinvals: ", self.sinvals[:10].cpu().numpy())
        self.inv_proj = V.transpose(0, 1)
        self.N = X.size(0)
//  [python]
def __init__(self, X, ndim=128, var_fraction=0.99, l2_normalized=True, first_direction=None):

        self.l2_normalized = l2_normalized
        if l2_normalized:
            X = X[:, :-1]

        assert len(X.shape) == 2
        torch.cuda.synchronize()
        start_time = time.time()

        self.mean = torch.mean(X, dim=0, keepdim=True)
        #self.mean = 0
        #self.std = torch.std(X, dim=0, keepdim=True) + 1e-6
        self.std = 1
        #print("std is ", self.std[:, :10].cpu().numpy())
        #X_orig = X
        X = (X - self.mean) / self.std

        U, S, V = torch.linalg.svd(X)
        S = S[:ndim]
        V = V[:, :ndim]
        self.proj = V
        scale = torch.mm(X, self.proj).std(dim=0)
        torch.cuda.synchronize()
        print("PCA time taken on vectors of size %s : %f" % (str(X.size()), time.time() - start_time))
        print("largest std of each PC: ", scale[:10].cpu().numpy())
        print("smallest std of each PC: ", scale[-10:].cpu().numpy())
        self.sinvals = S
        print("largest sinvals: ", self.sinvals[:10].cpu().numpy())
        self.inv_proj = V.transpose(0, 1)
        self.N = X.size(0)
//  [python]
def polar_decomp_left(M, normalize=False):
    U, S, V= torch.svd(M, compute_uv=True)
    # 
    # M= --U--S--V^T-- = --(U--V^T)--(V--S--V^T)-- = --U'--P--                                                             --/            2
    U= U@V.t()
    if normalize: S=S/S[0]
    P= V@torch.diag(S)@V.t()
    return P, U
//  [python]
def polar_decomp_left(M, normalize=False):
    U, S, V= torch.linalg.svd(M)
    U= U@V.t()
    if normalize: S=S/S[0]
    P= V@torch.diag(S)@V.t()
    return P, U
//  [python]
def get_matrix_kernel(A, eps=1e-10):
    '''
    Compute an orthonormal basis of the kernel (x_1, x_2, ...)
    A x_i = 0
    scalar_product(x_i, x_j) = delta_ij

    :param A: matrix
    :return: matrix where each row is a basis vector of the kernel of A
    '''
    _u, s, v = torch.svd(A)

    # A = u @ torch.diag(s) @ v.t()
    kernel = v.t()[s < eps]
    return kernel
//  [python]
def get_matrix_kernel(A, eps=1e-10):
    '''
    Compute an orthonormal basis of the kernel (x_1, x_2, ...)
    A x_i = 0
    scalar_product(x_i, x_j) = delta_ij

    :param A: matrix
    :return: matrix where each row is a basis vector of the kernel of A
    '''
    _u, s, v = torch.linalg.svd(A)

    # A = u @ torch.diag(s) @ v.t()
    kernel = v.t()[s < eps]
    return kernel
//  [python]
def superimpose(pred, true, atom_mask):
    
    def centroid(X):
        return X.mean(dim=-2, keepdim=True)
    
    B, L, natoms = pred.shape[:3]

    # center to centroid
    pred_allatom = pred[atom_mask][None]
    true_allatom = true[atom_mask][None]

    cp = centroid(pred_allatom)
    ct = centroid(true_allatom)
    
    pred_allatom_origin = pred_allatom - cp
    true_allatom_origin = true_allatom - ct

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

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

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

    # Rotation matrix U
    U = torch.matmul(d*V, W.permute(0,2,1)) # (IB, 3, 3)
    pred_rms = pred - cp
    true_rms = true - ct
    
    # Rotate pred
    rP = torch.matmul(pred_rms, U) # (IB, L*3, 3)
    
    return rP+ct
//  [python]
def superimpose(pred, true, atom_mask):
    
    def centroid(X):
        return X.mean(dim=-2, keepdim=True)
    
    B, L, natoms = pred.shape[:3]

    # center to centroid
    pred_allatom = pred[atom_mask][None]
    true_allatom = true[atom_mask][None]

    cp = centroid(pred_allatom)
    ct = centroid(true_allatom)
    
    pred_allatom_origin = pred_allatom - cp
    true_allatom_origin = true_allatom - ct

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

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

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

    # Rotation matrix U
    U = torch.matmul(d*U, V.permute(0,2,1)) # (IB, 3, 3)
    pred_rms = pred - cp
    true_rms = true - ct
    
    # Rotate pred
    rP = torch.matmul(pred_rms, U) # (IB, L*3, 3)
    
    return rP+ct
//  [python]
def __init__(self,encode_layers,Nkoopman,u_dim):
        super(Network,self).__init__()
        Layers = OrderedDict()
        for layer_i in range(len(encode_layers)-1):
            Layers["linear_{}".format(layer_i)] = nn.Linear(encode_layers[layer_i],encode_layers[layer_i+1])
            if layer_i != len(encode_layers)-2:
                Layers["relu_{}".format(layer_i)] = nn.ReLU()
        self.encode_net = nn.Sequential(Layers)
        self.Nkoopman = Nkoopman
        self.u_dim = u_dim
        self.lA = nn.Linear(Nkoopman,Nkoopman,bias=False)
        self.lA.weight.data = gaussian_init_(Nkoopman, std=1)
        U, _, V = torch.svd(self.lA.weight.data)
        self.lA.weight.data = torch.mm(U, V.t()) * 0.9
        self.lB = nn.Linear(u_dim,Nkoopman,bias=False)
//  [python]
def __init__(self,encode_layers,Nkoopman,u_dim):
        super(Network,self).__init__()
        Layers = OrderedDict()
        for layer_i in range(len(encode_layers)-1):
            Layers["linear_{}".format(layer_i)] = nn.Linear(encode_layers[layer_i],encode_layers[layer_i+1])
            if layer_i != len(encode_layers)-2:
                Layers["relu_{}".format(layer_i)] = nn.ReLU()
        self.encode_net = nn.Sequential(Layers)
        self.Nkoopman = Nkoopman
        self.u_dim = u_dim
        self.lA = nn.Linear(Nkoopman,Nkoopman,bias=False)
        self.lA.weight.data = gaussian_init_(Nkoopman, std=1)
        U, S, V = torch.linalg.svd(self.lA.weight.data)
        self.lA.weight.data = torch.mm(U, V.t()) * 0.9
        self.lB = nn.Linear(u_dim,Nkoopman,bias=False)
//  [python]
def pointcloud_rigid_registration(src_pointcloud, dst_pointcloud, reduce_loss: bool = True):
    """
    Calculate RT and residual L2 loss for two pointclouds
    :param src_pointcloud: x (b, v, 3)
    :param dst_pointcloud: y (b, v, 3)
    :return: loss, R, t  s.t. ||Rx+t-y||_2^2 minimal.
    """
    if len(src_pointcloud.shape) == 2:
        src_pointcloud = src_pointcloud.unsqueeze(0)
    if len(dst_pointcloud.shape) == 2:
        dst_pointcloud = dst_pointcloud.unsqueeze(0)
    bn = src_pointcloud.shape[0]

    assert src_pointcloud.shape == dst_pointcloud.shape
    assert src_pointcloud.shape[2] == 3

    X = src_pointcloud - src_pointcloud.mean(dim=1, keepdim=True)
    Y = dst_pointcloud - dst_pointcloud.mean(dim=1, keepdim=True)

    XYT = th.einsum("nji,njk->nik", X, Y)
    muX = src_pointcloud.mean(dim=1)
    muY = dst_pointcloud.mean(dim=1)

    R = th.zeros((bn, 3, 3), dtype=src_pointcloud.dtype).to(src_pointcloud.device)
    t = th.zeros((bn, 1, 3), dtype=src_pointcloud.dtype).to(src_pointcloud.device)
    loss = th.zeros((bn,), dtype=src_pointcloud.dtype).to(src_pointcloud.device)

    for i in range(bn):
        u_, s_, v_ = th.svd(XYT[i, :, :])
        detvut = th.det(v_.mm(u_.t()))
        diag_m = th.ones_like(s_)
        diag_m[-1] = detvut

        r_ = v_.mm(th.diag(diag_m)).mm(u_.t())
        t_ = muY[i, :] - r_.mm(muX[i, :, None])[:, 0]

        R[i, :, :] = r_
        t[i, 0, :] = t_
        loss[i] = (th.einsum("ij,nj->ni", r_, X[i]) - Y[i]).pow(2).sum(1).mean(0)

    loss = loss.mean(0) if reduce_loss else loss
    return loss, R, t
//  [python]
def pointcloud_rigid_registration(src_pointcloud, dst_pointcloud, reduce_loss: bool = True):
    """
    Calculate RT and residual L2 loss for two pointclouds
    :param src_pointcloud: x (b, v, 3)
    :param dst_pointcloud: y (b, v, 3)
    :return: loss, R, t  s.t. ||Rx+t-y||_2^2 minimal.
    """
    if len(src_pointcloud.shape) == 2:
        src_pointcloud = src_pointcloud.unsqueeze(0)
    if len(dst_pointcloud.shape) == 2:
        dst_pointcloud = dst_pointcloud.unsqueeze(0)
    bn = src_pointcloud.shape[0]

    assert src_pointcloud.shape == dst_pointcloud.shape
    assert src_pointcloud.shape[2] == 3

    X = src_pointcloud - src_pointcloud.mean(dim=1, keepdim=True)
    Y = dst_pointcloud - dst_pointcloud.mean(dim=1, keepdim=True)

    XYT = th.einsum("nji,njk->nik", X, Y)
    muX = src_pointcloud.mean(dim=1)
    muY = dst_pointcloud.mean(dim=1)

    R = th.zeros((bn, 3, 3), dtype=src_pointcloud.dtype).to(src_pointcloud.device)
    t = th.zeros((bn, 1, 3), dtype=src_pointcloud.dtype).to(src_pointcloud.device)
    loss = th.zeros((bn,), dtype=src_pointcloud.dtype).to(src_pointcloud.device)

    for i in range(bn):
        u_, s_, v_ = th.linalg.svd(XYT[i, :, :])
        detvut = th.det(v_.mm(u_.t()))
        diag_m = th.ones_like(s_)
        diag_m[-1] = detvut

        r_ = v_.mm(th.diag(diag_m)).mm(u_.t())
        t_ = muY[i, :] - r_.mm(muX[i, :, None])[:, 0]

        R[i, :, :] = r_
        t[i, 0, :] = t_
        loss[i] = (th.einsum("ij,nj->ni", r_, X[i]) - Y[i]).pow(2).sum(1).mean(0)

    loss = loss.mean(0) if reduce_loss else loss
    return loss, R, t
//  [python]
def _symmetric_matrix_square_root(mat, eps=1e-10):
    u, s, v = torch.svd(mat)
    si = torch.where(s < eps, s, torch.sqrt(s))
    return torch.matmul(torch.matmul(u, torch.diag(si)), v.t())
//  [python]
def _symmetric_matrix_square_root(mat, eps=1e-10):
    u, s, v = torch.linalg.svd(mat)
    si = torch.where(s < eps, s, torch.sqrt(s))
    return torch.matmul(torch.matmul(u, torch.diag(si)), v.t())
//  [python]
def pca(X, k=2):
    X_mean = torch.mean(X,0)
    X = X - X_mean.expand_as(X)
    U,S,V = torch.svd(torch.t(X))
    return torch.mm(X,U[:,:k])
//  [python]
def pca(X, k=2):
    X_mean = torch.mean(X,0)
    X = X - X_mean.expand_as(X)
    U,S,V = torch.linalg.svd(torch.t(X))
    return torch.mm(X,U[:,:k])
//  [python]
def regularizer_orth(m):
    """
    # ----------------------------------------
    # SVD Orthogonal Regularization
    # ----------------------------------------
    # Applies regularization to the training by performing the
    # orthogonalization technique described in the paper
    # This function is to be called by the torch.nn.Module.apply() method,
    # which applies svd_orthogonalization() to every layer of the model.
    # usage: net.apply(regularizer_orth)
    # ----------------------------------------
    """
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        w = m.weight.data.clone()
        c_out, c_in, f1, f2 = w.size()
        # dtype = m.weight.data.type()
        w = w.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)
        # self.netG.apply(svd_orthogonalization)
        u, s, v = torch.svd(w)
        s[s > 1.5] = s[s > 1.5] - 1e-4
        s[s < 0.5] = s[s < 0.5] + 1e-4
        w = torch.mm(torch.mm(u, torch.diag(s)), v.t())
        m.weight.data = w.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1)  # .type(dtype)
    else:
        pass
//  [python]
def regularizer_orth(m):
    """
    # ----------------------------------------
    # SVD Orthogonal Regularization
    # ----------------------------------------
    # Applies regularization to the training by performing the
    # orthogonalization technique described in the paper
    # This function is to be called by the torch.nn.Module.apply() method,
    # which applies svd_orthogonalization() to every layer of the model.
    # usage: net.apply(regularizer_orth)
    # ----------------------------------------
    """
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        w = m.weight.data.clone()
        c_out, c_in, f1, f2 = w.size()
        # dtype = m.weight.data.type()
        w = w.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)
        # self.netG.apply(svd_orthogonalization)
        u, s, v = torch.linalg.svd(w)
        s[s > 1.5] = s[s > 1.5] - 1e-4
        s[s < 0.5] = s[s < 0.5] + 1e-4
        w = torch.mm(torch.mm(u, torch.diag(s)), v.t())
        m.weight.data = w.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1)  # .type(dtype)
    else:
        pass
//  [python]
def regularizer_orth2(m):
    """
    # ----------------------------------------
    # Applies regularization to the training by performing the
    # orthogonalization technique described in the paper
    # This function is to be called by the torch.nn.Module.apply() method,
    # which applies svd_orthogonalization() to every layer of the model.
    # usage: net.apply(regularizer_orth2)
    # ----------------------------------------
    """
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        w = m.weight.data.clone()
        c_out, c_in, f1, f2 = w.size()
        # dtype = m.weight.data.type()
        w = w.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)
        u, s, v = torch.svd(w)
        s_mean = s.mean()
        s[s > 1.5*s_mean] = s[s > 1.5*s_mean] - 1e-4
        s[s < 0.5*s_mean] = s[s < 0.5*s_mean] + 1e-4
        w = torch.mm(torch.mm(u, torch.diag(s)), v.t())
        m.weight.data = w.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1)  # .type(dtype)
    else:
        pass
//  [python]
def regularizer_orth2(m):
    """
    # ----------------------------------------
    # Applies regularization to the training by performing the
    # orthogonalization technique described in the paper
    # This function is to be called by the torch.nn.Module.apply() method,
    # which applies svd_orthogonalization() to every layer of the model.
    # usage: net.apply(regularizer_orth2)
    # ----------------------------------------
    """
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        w = m.weight.data.clone()
        c_out, c_in, f1, f2 = w.size()
        # dtype = m.weight.data.type()
        w = w.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)
        u, s, v = torch.linalg.svd(w)
        s_mean = s.mean()
        s[s > 1.5*s_mean] = s[s > 1.5*s_mean] - 1e-4
        s[s < 0.5*s_mean] = s[s < 0.5*s_mean] + 1e-4
        w = torch.mm(torch.mm(u, torch.diag(s)), v.t())
        m.weight.data = w.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1)  # .type(dtype)
    else:
        pass
//  [python]
def svd(
        tensor: torch.Tensor,
        reg_value: float = 0.0
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Performs singular value decomposition of a tensor. The regularization
    value is added to the tensor.

    Parameters
    ----------
    tensor : :py:class:`torch.Tensor`
        This tensor is decomposed.
    reg_value : float, optional
        This regularization value is added to the singular values and represents
        a regularization of the matrix diagonal. The regularization is
        deactivated with the default value of 0.

    Returns
    -------
    u : :py:class:`torch.Tensor`
        The decomposed left singular vector.
    s : :py:class:`torch.Tensor`
        The decomposed singular values with added regularization value.
    v : :py:class:`torch.Tensor`
        The decomposed right singular vector.
    """
    u, s, v = torch.svd(tensor)
    s = s + reg_value
    return u, s, v
//  [python]
def svd(
        tensor: torch.Tensor,
        reg_value: float = 0.0
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Performs singular value decomposition of a tensor. The regularization
    value is added to the tensor.

    Parameters
    ----------
    tensor : :py:class:`torch.Tensor`
        This tensor is decomposed.
    reg_value : float, optional
        This regularization value is added to the singular values and represents
        a regularization of the matrix diagonal. The regularization is
        deactivated with the default value of 0.

    Returns
    -------
    u : :py:class:`torch.Tensor`
        The decomposed left singular vector.
    s : :py:class:`torch.Tensor`
        The decomposed singular values with added regularization value.
    v : :py:class:`torch.Tensor`
        The decomposed right singular vector.
    """
    u, s, v = torch.linalg.svd(tensor)
    s = s + reg_value
    return u, s, v
//  [python]
def ex_net_svd(model, in_dim):
  """
  Performs a Singular Value Decomposition on a given model weights

  Args:
    model: torch.nn.Module
      Neural network model
    in_dim: int
      The input dimension of the model

  Returns:
    U: torch.tensor
      Orthogonal matrix
    Σ: torch.tensor
      Diagonal matrix
    V: torch.tensor
      Orthogonal matrix
  """
  W_tot = torch.eye(in_dim)
  for weight in model.parameters():
    W_tot = weight @ W_tot
  U, Σ, V = torch.svd(W_tot)
  return U, Σ, V
//  [python]
def ex_net_svd(model, in_dim):
  """
  Performs a Singular Value Decomposition on a given model weights

  Args:
    model: torch.nn.Module
      Neural network model
    in_dim: int
      The input dimension of the model

  Returns:
    U: torch.tensor
      Orthogonal matrix
    Σ: torch.tensor
      Diagonal matrix
    V: torch.tensor
      Orthogonal matrix
  """
  W_tot = torch.eye(in_dim)
  for weight in model.parameters():
    W_tot = weight @ W_tot
  U, S, V = torch.linalg.svd(W_tot)
  return U, S, V
//  [python]
def decompose_essential_matrix(E):
    # decompose matrix by its singular values
    U, _, V = torch.svd(E)
    Vt = V.transpose(-2, -1)

    mask = torch.ones_like(E)
    mask[..., -1:] *= -1.0  # fill last column with negative values

    maskt = mask.transpose(-2, -1)

    # avoid singularities
    U = torch.where((torch.det(U) < 0.0)[..., None, None], U * mask, U)
    Vt = torch.where((torch.det(Vt) < 0.0)[..., None, None], Vt * maskt, Vt)

    W = skew_symmetric(E.new_tensor([[0, 0, 1]]))
    W[..., 2, 2] += 1.0

    # reconstruct rotations and retrieve translation vector
    U_W_Vt = U @ W @ Vt
    U_Wt_Vt = U @ W.transpose(-2, -1) @ Vt

    # return values
    R1 = U_W_Vt
    R2 = U_Wt_Vt
    T = U[..., -1]
    return R1, R2, T
//  [python]
def decompose_essential_matrix(E):
    # decompose matrix by its singular values
    U, S, V = torch.linalg.svd(E)
    Vt = V.transpose(-2, -1)

    mask = torch.ones_like(E)
    mask[..., -1:] *= -1.0  # fill last column with negative values

    maskt = mask.transpose(-2, -1)

    # avoid singularities
    U = torch.where((torch.det(U) < 0.0)[..., None, None], U * mask, U)
    Vt = torch.where((torch.det(Vt) < 0.0)[..., None, None], Vt * maskt, Vt)

    W = skew_symmetric(E.new_tensor([[0, 0, 1]]))
    W[..., 2, 2] += 1.0

    # reconstruct rotations and retrieve translation vector
    U_W_Vt = U @ W @ Vt
    U_Wt_Vt = U @ W.transpose(-2, -1) @ Vt

    # return values
    R1 = U_W_Vt
    R2 = U_Wt_Vt
    T = U[..., -1]
    return R1, R2, T
//  [python]
def __init__(self, in_features: int, out_features: int, weight_std: float, bias_std: float):
        super(INRSVDLinear, self).__init__(in_features, out_features, weight_std=weight_std, bias_std=bias_std)

        weight = torch.empty(out_features, in_features)
        bound = weight_std * np.sqrt(3)
        nn.init.uniform_(weight, -bound, bound)

        U, _, V = torch.svd(weight)
        self.U = nn.Parameter(U)
        self.V = nn.Parameter(V)

        self.in_features = in_features
        self.out_features = out_features
        self.num_singular_values = min(in_features, out_features)
        self.num_external_params = self.num_singular_values + out_features
//  [python]
def __init__(self, in_features: int, out_features: int, weight_std: float, bias_std: float):
        super(INRSVDLinear, self).__init__(in_features, out_features, weight_std=weight_std, bias_std=bias_std)

        weight = torch.empty(out_features, in_features)
        bound = weight_std * np.sqrt(3)
        nn.init.uniform_(weight, -bound, bound)

        U, S, V = torch.linalg.svd(weight)
        self.U = nn.Parameter(U)
        self.V = nn.Parameter(V)

        self.in_features = in_features
        self.out_features = out_features
        self.num_singular_values = min(in_features, out_features)
        self.num_external_params = self.num_singular_values + out_features
//  [python]
def test_svd_complex(self, device, dtype):
        t = torch.randn((10, 10), dtype=dtype, device=device)
        U, S, V = torch.svd(t, some=False)
        # note: from the math point of view, it is weird that we need to use
        # V.T instead of V.T.conj(): torch.svd has a buggy behavior for
        # complex numbers and it's deprecated. You should use torch.linalg.svd
        # instead.
        t2 = U @ torch.diag(S).type(dtype) @ V.T
        self.assertEqual(t, t2)
//  [python]
def test_svd_complex(self, device, dtype):
        t = torch.randn((10, 10), dtype=dtype, device=device)
        U, S, V = torch.linalg.svd(t)
        t2 = U @ torch.diag(S).type(dtype) @ V.T
        self.assertEqual(t, t2)
//  [python]
def test_svd_no_singularvectors(self, device, dtype):
        for size in [(5, 5), (5, 20), (20, 5)]:
            a = torch.randn(*size, device=device, dtype=dtype)
            u, s_expect, v = torch.svd(a)
            u, s_actual, v = torch.svd(a, compute_uv=False)
            self.assertEqual(s_expect, s_actual, msg="Singular values don't match")
//  [python]
def test_svd_no_singularvectors(self, device, dtype):
        for size in [(5, 5), (5, 20), (20, 5)]:
            a = torch.randn(*size, device=device, dtype=dtype)
            u, s_expect, v = torch.linalg.svd(a)
            u, s_actual, v = torch.linalg.svd(a, compute_uv=False)
            self.assertEqual(s_expect, s_actual, msg="Singular values don't match")
//  [python]
def kabsch_umeyama(A, B):
    n, m = A.shape
    EA = torch.mean(A, axis=0)
    EB = torch.mean(B, axis=0)
    VarA = torch.mean((A - EA).norm(dim=1)**2)

    H = ((A - EA).T @ (B - EB)) / n
    U, D, VT = torch.svd(H)

    c = VarA / torch.trace(torch.diag(D))
    return c
//  [python]
def kabsch_umeyama(A, B):
    n, m = A.shape
    EA = torch.mean(A, axis=0)
    EB = torch.mean(B, axis=0)
    VarA = torch.mean((A - EA).norm(dim=1)**2)

    H = ((A - EA).T @ (B - EB)) / n
    U, S, VT = torch.linalg.svd(H)

    c = VarA / torch.trace(torch.diag(S))
    return c
//  [python]
def _test_svd_helper(self, shape, some, col_maj, device, dtype):
        cpu_tensor = torch.randn(shape, device='cpu').to(dtype)
        device_tensor = cpu_tensor.to(device=device)
        if col_maj:
            cpu_tensor = cpu_tensor.t()
            device_tensor = device_tensor.t()
        cpu_result = torch.svd(cpu_tensor, some=some)
        device_result = torch.svd(device_tensor, some=some)
        m = min(cpu_tensor.shape[-2:])
        # torch.svd returns torch.return_types.svd which is a tuple of (U, V, S).
        # - When some==False, U[..., m:] can be arbitrary.
        # - When some==True, U shape: [..., m], V shape: [m, m]
        # - Signs are not deterministic. If the sign of a column of U is changed
        #   then the corresponding column of the V has to be changed.
        # Thus here we only compare result[..., :m].abs() from CPU and device.
        for x, y in zip(cpu_result, device_result):
            self.assertEqual(x[..., :m].abs(), y[..., :m].abs(), atol=1e-5, rtol=0)
//  [python]
def _test_svd_helper(self, shape, some, col_maj, device, dtype):
        cpu_tensor = torch.randn(shape, device='cpu').to(dtype)
        device_tensor = cpu_tensor.to(device=device)
        if col_maj:
            cpu_tensor = cpu_tensor.t()
            device_tensor = device_tensor.t()
        cpu_result = torch.linalg.svd(cpu_tensor, some=some)
        device_result = torch.linalg.svd(device_tensor, some=some)
        m = min(cpu_tensor.shape[-2:])
        # torch.linalg.svd returns a tuple of (U, S, Vh).
        # - When some==False, U[..., m:] can be arbitrary.
        # - When some==True, U shape: [..., m], Vh shape: [m, m]
        # - Signs are not deterministic. If the sign of a column of U is changed
        #   then the corresponding column of the V has to be changed.
        # Thus here we only compare result[..., :m].abs() from CPU and device.
        for x, y in zip(cpu_result, device_result):
            self.assertEqual(x[..., :m].abs(), y[..., :m].abs(), atol=1e-5, rtol=0)
//  [python]
def get_regularized_kernel(self, kernel):
        if self.args.kernel_regularization:
            _, S, _ = torch.svd(kernel)
            op_norm = S.max()

            reg = (self.args.kernel_regularization * op_norm).to(kernel.dtype)

            identity = torch.eye(kernel.shape[1], dtype=kernel.dtype).unsqueeze(0)
            return kernel + identity * reg
        else:
            return kernel
//  [python]
def get_regularized_kernel(self, kernel):
        if self.args.kernel_regularization:
            U, S, V = torch.linalg.svd(kernel)
            op_norm = S.max()

            reg = (self.args.kernel_regularization * op_norm).to(kernel.dtype)

            identity = torch.eye(kernel.shape[1], dtype=kernel.dtype).unsqueeze(0)
            return kernel + identity * reg
        else:
            return kernel
//  [python]
def get_jacobian(self, func, inputs, device, create_graph=False, strict=False):
        for batch_idx, (data, target) in enumerate(inputs):
            if batch_idx ==0:
                print('data len:', len(data))
                data, target = data.to(device), target.to(device)
                jacobian_4d = torch.autograd.functional.jacobian(func, data, create_graph=create_graph,
                                                                 strict=strict)  # 100, 10, 100, 784
                jacobian_4d = jacobian_4d.view(self.args.batch_size, 10, self.args.batch_size, -1)
                jacobian_3d = torch.zeros(jacobian_4d.size(0), jacobian_4d.size(1), jacobian_4d.size(3))  # 100, 10, 784
                for i in range(jacobian_3d.size(0)):
                    for j in range(jacobian_3d.size(1)):
                        jacobian_3d[i][j] = jacobian_4d[i][j][i]
                # u, s, vh = torch.svd(jacobian_3d.cpu().numpy()[1], full_matrices=True)
                svj = torch.svd(jacobian_3d)[1]
                print('SVJ is:', svj.mean(), 'SJV size is:', svj.size())
                return svj
            else: break
//  [python]
def get_jacobian(self, func, inputs, device, create_graph=False, strict=False):
        for batch_idx, (data, target) in enumerate(inputs):
            if batch_idx ==0:
                print('data len:', len(data))
                data, target = data.to(device), target.to(device)
                jacobian_4d = torch.autograd.functional.jacobian(func, data, create_graph=create_graph,
                                                                 strict=strict)  # 100, 10, 100, 784
                jacobian_4d = jacobian_4d.view(self.args.batch_size, 10, self.args.batch_size, -1)
                jacobian_3d = torch.zeros(jacobian_4d.size(0), jacobian_4d.size(1), jacobian_4d.size(3))  # 100, 10, 784
                for i in range(jacobian_3d.size(0)):
                    for j in range(jacobian_3d.size(1)):
                        jacobian_3d[i][j] = jacobian_4d[i][j][i]
                # u, s, vh = torch.svd(jacobian_3d.cpu().numpy()[1], full_matrices=True)
                svj = torch.linalg.svd(jacobian_3d)[1]
                print('SVJ is:', svj.mean(), 'SJV size is:', svj.size())
                return svj
            else: break
//  [python]
def torch_align_w_scale(mtx1, mtx2, return_trafo=False):
    """ Align the predicted entity in some optimality sense with the ground truth. """
    # center
    t1 = torch.mean(mtx1, dim=1, keepdim=True)
    t2 = torch.mean(mtx2, dim=1, keepdim=True)
    mtx1_t = mtx1 - t1
    mtx2_t = mtx2 - t2

    # scale
    s1 = torch.linalg.matrix_norm(mtx1_t, dim=(-2, -1), keepdim=True) + 1e-8
    mtx1_t /= s1
    s2 = torch.linalg.matrix_norm(mtx2_t, dim=(-2, -1), keepdim=True) + 1e-8
    mtx2_t /= s2

    # orth alignment
    u, w, v = torch.svd(torch.matmul(mtx2_t.transpose(1, 2), mtx1_t).transpose(1, 2))
    R = torch.matmul(u, v.transpose(1, 2))
    s = torch.sum(w, dim=1, keepdim=True).unsqueeze(-1)

    # apply trafos to the second matrix
    mtx2_t = torch.matmul(mtx2_t, R.transpose(1, 2)) * s
    mtx2_t = mtx2_t * s1 + t1

    if return_trafo:
        return R, s, s1, t1 - t2
    else:
        return mtx2_t
//  [python]
def torch_align_w_scale(mtx1, mtx2, return_trafo=False):
    """ Align the predicted entity in some optimality sense with the ground truth. """
    # center
    t1 = torch.mean(mtx1, dim=1, keepdim=True)
    t2 = torch.mean(mtx2, dim=1, keepdim=True)
    mtx1_t = mtx1 - t1
    mtx2_t = mtx2 - t2

    # scale
    s1 = torch.linalg.matrix_norm(mtx1_t, dim=(-2, -1), keepdim=True) + 1e-8
    mtx1_t /= s1
    s2 = torch.linalg.matrix_norm(mtx2_t, dim=(-2, -1), keepdim=True) + 1e-8
    mtx2_t /= s2

    # orth alignment
    u, w, v = torch.linalg.svd(torch.matmul(mtx2_t.transpose(1, 2), mtx1_t).transpose(1, 2))
    R = torch.matmul(u, v.transpose(1, 2))
    s = torch.sum(w, dim=1, keepdim=True).unsqueeze(-1)

    # apply trafos to the second matrix
    mtx2_t = torch.matmul(mtx2_t, R.transpose(1, 2)) * s
    mtx2_t = mtx2_t * s1 + t1

    if return_trafo:
        return R, s, s1, t1 - t2
    else:
        return mtx2_t
//  [python]
def get_cameras_accuracy(pred_Rs, gt_Rs, pred_ts, gt_ts,):
    ''' Align predicted pose to gt pose and print cameras accuracy'''

    # find rotation
    d = pred_Rs.shape[-1]
    n = pred_Rs.shape[0]

    Q = torch.addbmm(torch.zeros(d, d, dtype=torch.double), gt_Rs, pred_Rs.transpose(1, 2))
    Uq, _, Vq = torch.svd(Q)
    sv = torch.ones(d, dtype=torch.double)
    sv[-1] = torch.det(Uq @ Vq.transpose(0, 1))
    R_opt = Uq @ torch.diag(sv) @ Vq.transpose(0, 1)
    R_fixed = torch.bmm(R_opt.repeat(n, 1, 1), pred_Rs)

    # find translation
    pred_ts = pred_ts @ R_opt.transpose(0, 1)
    c_opt = cp.Variable()
    t_opt = cp.Variable((1, d))

    constraints = []
    obj = cp.Minimize(cp.sum(
        cp.norm(gt_ts.numpy() - (c_opt * pred_ts.numpy() + np.ones((n, 1), dtype=np.double) @ t_opt), axis=1)))
    prob = cp.Problem(obj, constraints)
    prob.solve()
    t_fixed = c_opt.value * pred_ts.numpy() + np.ones((n, 1), dtype=np.double) * t_opt.value

    # Calculate transaltion error
    t_error = np.linalg.norm(t_fixed - gt_ts.numpy(), axis=-1)
    t_error = t_error
    t_error_mean = np.mean(t_error)
    t_error_medi = np.median(t_error)

    # Calculate rotation error
    R_error = compare_rotations(R_fixed, gt_Rs)

    R_error = R_error.numpy()
    R_error_mean = np.mean(R_error)
    R_error_medi = np.median(R_error)

    print('CAMERAS EVALUATION: R error mean = {0} ; t error mean = {1} ; R error median = {2} ; t error median = {3}'
          .format("%.2f" % R_error_mean, "%.2f" % t_error_mean, "%.2f" % R_error_medi, "%.2f" % t_error_medi))

    # return alignment and aligned pose
    return R_opt.numpy(), t_opt.value, c_opt.value, R_fixed.numpy(), t_fixed
//  [python]
def get_cameras_accuracy(pred_Rs, gt_Rs, pred_ts, gt_ts,):
    ''' Align predicted pose to gt pose and print cameras accuracy'''

    # find rotation
    d = pred_Rs.shape[-1]
    n = pred_Rs.shape[0]

    Q = torch.addbmm(torch.zeros(d, d, dtype=torch.double), gt_Rs, pred_Rs.transpose(1, 2))
    Uq, S, Vq = torch.linalg.svd(Q)
    sv = torch.ones(d, dtype=torch.double)
    sv[-1] = torch.det(Uq @ Vq.transpose(0, 1))
    R_opt = Uq @ torch.diag(sv) @ Vq.transpose(0, 1)
    R_fixed = torch.bmm(R_opt.repeat(n, 1, 1), pred_Rs)

    # find translation
    pred_ts = pred_ts @ R_opt.transpose(0, 1)
    c_opt = cp.Variable()
    t_opt = cp.Variable((1, d))

    constraints = []
    obj = cp.Minimize(cp.sum(
        cp.norm(gt_ts.numpy() - (c_opt * pred_ts.numpy() + np.ones((n, 1), dtype=np.double) @ t_opt), axis=1)))
    prob = cp.Problem(obj, constraints)
    prob.solve()
    t_fixed = c_opt.value * pred_ts.numpy() + np.ones((n, 1), dtype=np.double) * t_opt.value

    # Calculate transaltion error
    t_error = np.linalg.norm(t_fixed - gt_ts.numpy(), axis=-1)
    t_error = t_error
    t_error_mean = np.mean(t_error)
    t_error_medi = np.median(t_error)

    # Calculate rotation error
    R_error = compare_rotations(R_fixed, gt_Rs)

    R_error = R_error.numpy()
    R_error_mean = np.mean(R_error)
    R_error_medi = np.median(R_error)

    print('CAMERAS EVALUATION: R error mean = {0} ; t error mean = {1} ; R error median = {2} ; t error median = {3}'
          .format("%.2f" % R_error_mean, "%.2f" % t_error_mean, "%.2f" % R_error_medi, "%.2f" % t_error_medi))

    # return alignment and aligned pose
    return R_opt.numpy(), t_opt.value, c_opt.value, R_fixed.numpy(), t_fixed
//  [python]
def test_svd_lowrank(self, device, dtype):
        from torch.testing._internal.common_utils import random_lowrank_matrix, random_sparse_matrix

        def run_subtest(actual_rank, matrix_size, batches, device, svd_lowrank, **options):
            density = options.pop('density', 1)
            if isinstance(matrix_size, int):
                rows = columns = matrix_size
            else:
                rows, columns = matrix_size
            if density == 1:
                a_input = random_lowrank_matrix(actual_rank, rows, columns, *batches, device=device, dtype=dtype)
                a = a_input
            else:
                assert batches == ()
                a_input = random_sparse_matrix(rows, columns, density, device=device, dtype=dtype)
                a = a_input.to_dense()

            q = min(*size)
            u, s, v = svd_lowrank(a_input, q=q, **options)

            # check if u, s, v is a SVD
            u, s, v = u[..., :q], s[..., :q], v[..., :q]
            A = u.matmul(s.diag_embed()).matmul(v.transpose(-2, -1))
            self.assertEqual(A, a)

            # check if svd_lowrank produces same singular values as torch.svd
            U, S, V = torch.svd(a)
            self.assertEqual(s.shape, S.shape)
            self.assertEqual(u.shape, U.shape)
            self.assertEqual(v.shape, V.shape)
            self.assertEqual(s, S)

            if density == 1:
                # actual_rank is known only for dense inputs
                #
                # check if pairs (u, U) and (v, V) span the same
                # subspaces, respectively
                u, s, v = u[..., :actual_rank], s[..., :actual_rank], v[..., :actual_rank]
                U, S, V = U[..., :actual_rank], S[..., :actual_rank], V[..., :actual_rank]
                self.assertEqual(u.transpose(-2, -1).matmul(U).det().abs(), torch.ones(batches, device=device, dtype=dtype))
                self.assertEqual(v.transpose(-2, -1).matmul(V).det().abs(), torch.ones(batches, device=device, dtype=dtype))

        all_batches = [(), (1,), (3,), (2, 3)]
        for actual_rank, size, all_batches in [
                (2, (17, 4), all_batches),
                (4, (17, 4), all_batches),
                (4, (17, 17), all_batches),
                (10, (100, 40), all_batches),
                (7, (1000, 1000), [()]),
        ]:
            # dense input
            for batches in all_batches:
                run_subtest(actual_rank, size, batches, device, torch.svd_lowrank)
                if size != size[::-1]:
                    run_subtest(actual_rank, size[::-1], batches, device, torch.svd_lowrank)

        # sparse input
        for size in [(17, 4), (4, 17), (17, 17), (100, 40), (40, 100), (1000, 1000)]:
            for density in [0.005, 0.1]:
                run_subtest(None, size, (), device, torch.svd_lowrank, density=density)

        # jitting support
        jitted = torch.jit.script(torch.svd_lowrank)
        actual_rank, size, batches = 2, (17, 4), ()
        run_subtest(actual_rank, size, batches, device, jitted)
//  [python]
def test_svd_lowrank(self, device, dtype):
        from torch.testing._internal.common_utils import random_lowrank_matrix, random_sparse_matrix

        def run_subtest(actual_rank, matrix_size, batches, device, svd_lowrank, **options):
            density = options.pop('density', 1)
            if isinstance(matrix_size, int):
                rows = columns = matrix_size
            else:
                rows, columns = matrix_size
            if density == 1:
                a_input = random_lowrank_matrix(actual_rank, rows, columns, *batches, device=device, dtype=dtype)
                a = a_input
            else:
                assert batches == ()
                a_input = random_sparse_matrix(rows, columns, density, device=device, dtype=dtype)
                a = a_input.to_dense()

            q = min(*size)
            u, s, v = svd_lowrank(a_input, q=q, **options)

            # check if u, s, v is a SVD
            u, s, v = u[..., :q], s[..., :q], v[..., :q]
            A = u.matmul(s.diag_embed()).matmul(v.transpose(-2, -1))
            self.assertEqual(A, a)

            # check if svd_lowrank produces same singular values as torch.linalg.svd
            U, S, V = torch.linalg.svd(a)
            self.assertEqual(s.shape, S.shape)
            self.assertEqual(u.shape, U.shape)
            self.assertEqual(v.shape, V.shape)
            self.assertEqual(s, S)

            if density == 1:
                # actual_rank is known only for dense inputs
                #
                # check if pairs (u, U) and (v, V) span the same
                # subspaces, respectively
                u, s, v = u[..., :actual_rank], s[..., :actual_rank], v[..., :actual_rank]
                U, S, V = U[..., :actual_rank], S[..., :actual_rank], V[..., :actual_rank]
                self.assertEqual(u.transpose(-2, -1).matmul(U).det().abs(), torch.ones(batches, device=device, dtype=dtype))
                self.assertEqual(v.transpose(-2, -1).matmul(V).det().abs(), torch.ones(batches, device=device, dtype=dtype))

        all_batches = [(), (1,), (3,), (2, 3)]
        for actual_rank, size, all_batches in [
                (2, (17, 4), all_batches),
                (4, (17, 4), all_batches),
                (4, (17, 17), all_batches),
                (10, (100, 40), all_batches),
                (7, (1000, 1000), [()]),
        ]:
            # dense input
            for batches in all_batches:
                run_subtest(actual_rank, size, batches, device, torch.svd_lowrank)
                if size != size[::-1]:
                    run_subtest(actual_rank, size[::-1], batches, device, torch.svd_lowrank)

        # sparse input
        for size in [(17, 4), (4, 17), (17, 17), (100, 40), (40, 100), (1000, 1000)]:
            for density in [0.005, 0.1]:
                run_subtest(None, size, (), device, torch.svd_lowrank, density=density)

        # jitting support
        jitted = torch.jit.script(torch.svd_lowrank)
        actual_rank, size, batches = 2, (17, 4), ()
        run_subtest(actual_rank, size, batches, device, jitted)
//  [python]
def procrustes(A, B):
    """
    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 (B, N, 3) -- Point Cloud to Align (source)
        -    B: Torch tensor of shape (B, N, 3) -- Reference Point Cloud (target)
        Returns:
        -    R: optimal rotation (B, 3, 3)
        -    t: optimal translation  (B, 3, 1)
    Based on: https://gist.github.com/bougui505/e392a371f5bab095a3673ea6f4976cc8
    """
    assert len(A.shape) == len(B.shape) == 3, 'three dimensions are required'
    assert A.shape[0] == B.shape[0], 'batch size must match'
    assert A.shape[1] == B.shape[1], 'number of correspondences must match'
    assert A.shape[2] == B.shape[2], 'number of spatial dimensions must be 3'

    a_mean = A.mean(axis=1, keepdim=True)
    b_mean = B.mean(axis=1, keepdim=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = A_c.transpose(1, 2) @ B_c
    U, S, V = torch.svd(H)
    # Fixes orientation such that Det(R) = + 1
    Z = torch.eye(3).unsqueeze(0).repeat(A.shape[0], 1, 1).to(A.device)
    Z[:, -1, -1] = torch.sign(torch.linalg.det(U @ V.transpose(1, 2)))
    # Rotation matrix
    R = V @ Z @ U.transpose(1, 2)
    # Translation vector
    t = b_mean - a_mean @ R.transpose(1, 2)
    return R, t
//  [python]
def procrustes(A, B):
    """
    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 (B, N, 3) -- Point Cloud to Align (source)
        -    B: Torch tensor of shape (B, N, 3) -- Reference Point Cloud (target)
        Returns:
        -    R: optimal rotation (B, 3, 3)
        -    t: optimal translation  (B, 3, 1)
    Based on: https://gist.github.com/bougui505/e392a371f5bab095a3673ea6f4976cc8
    """
    assert len(A.shape) == len(B.shape) == 3, 'three dimensions are required'
    assert A.shape[0] == B.shape[0], 'batch size must match'
    assert A.shape[1] == B.shape[1], 'number of correspondences must match'
    assert A.shape[2] == B.shape[2], 'number of spatial dimensions must be 3'

    a_mean = A.mean(axis=1, keepdim=True)
    b_mean = B.mean(axis=1, keepdim=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = A_c.transpose(1, 2) @ B_c
    U, S, V = torch.linalg.svd(H)
    # Fixes orientation such that Det(R) = + 1
    Z = torch.eye(3).unsqueeze(0).repeat(A.shape[0], 1, 1).to(A.device)
    Z[:, -1, -1] = torch.sign(torch.linalg.det(U @ V.transpose(1, 2)))
    # Rotation matrix
    R = V @ Z @ U.transpose(1, 2)
    # Translation vector
    t = b_mean - a_mean @ R.transpose(1, 2)
    return R, t
//  [python]
def compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.T
        S2 = S2.T
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=1, keepdims=True)
    mu2 = S2.mean(axis=1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # print('X1', X1.shape)

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

    # print('var', var1.shape)

    # 3. The outer product of X1 and X2.
    K = X1.mm(X2.T)

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, V = torch.svd(K)
    # V = Vh.T
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[0], device=S1.device)
    Z[-1, -1] *= torch.sign(torch.det(U @ V.T))
    # Construct R.
    R = V.mm(Z.mm(U.T))

    # print('R', X1.shape)

    # 5. Recover scale.
    scale = torch.trace(R.mm(K)) / var1
    # print(R.shape, mu1.shape)
    # 6. Recover translation.
    t = mu2 - scale * (R.mm(mu1))
    # print(t.shape)

    # 7. Error:
    S1_hat = scale * R.mm(S1) + t

    if transposed:
        S1_hat = S1_hat.T

    return S1_hat
//  [python]
def compute_similarity_transform_torch(S1, S2):
    '''
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    '''
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.T
        S2 = S2.T
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=1, keepdims=True)
    mu2 = S2.mean(axis=1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # print('X1', X1.shape)

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

    # print('var', var1.shape)

    # 3. The outer product of X1 and X2.
    K = X1.mm(X2.T)

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, S, V = torch.linalg.svd(K)
    # V = Vh.T
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[0], device=S1.device)
    Z[-1, -1] *= torch.sign(torch.det(U @ V.T))
    # Construct R.
    R = V.mm(Z.mm(U.T))

    # print('R', X1.shape)

    # 5. Recover scale.
    scale = torch.trace(R.mm(K)) / var1
    # print(R.shape, mu1.shape)
    # 6. Recover translation.
    t = mu2 - scale * (R.mm(mu1))
    # print(t.shape)

    # 7. Error:
    S1_hat = scale * R.mm(S1) + t

    if transposed:
        S1_hat = S1_hat.T

    return S1_hat
//  [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.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)