jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        min_cost = np.inf


        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...]
            l2_dis = torch.linalg.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), ord=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.linalg.norm(x[l2_cls == c] - tmp_center[c], ord=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.linalg.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), ord=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.linalg.norm((center_old - center), ord=2, dim=1).max()

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)
//  [python]
def init_prod(state, env, verbosity=0):
    for key,t in env.C.items():
        env.C[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.C[key][0,0]= 1.0 + 0.j if env.C[key].is_complex() else 1.0

    A= next(iter(state.sites.values()))
    # left transfer matrix
    #
    #     0      = 0     
    # i--A--j      T
    #   /\         1
    #  2  m
    #      \ 0
    #    i--A--j
    #      /
    #     2
    if len(A.size())==4:
        # assume only virtual indices are present
        a= torch.einsum('aibj->ab',A).contiguous()
    else:    
        a = torch.einsum('meifj,maibj->eafb',(A,A.conj())).contiguous().view(\
            A.size()[1]**2, A.size()[3]**2)
    a= a/a.abs().max()
    # check symmetry
    a_asymm_norm= torch.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, 2)
    # leading eigenvector is unique 
    assert torch.abs(D[0]-D[1]) > 1.0e-8, "Leading eigenvector of T not unique"
    for key,t in env.T.items():
        env.T[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.T[key][0,0,:]= U[:,0]
//  [python]
def init_prod(state, env, verbosity=0):
    for key,t in env.C.items():
        env.C[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.C[key][0,0]= 1.0 + 0.j if env.C[key].is_complex() else 1.0

    A= next(iter(state.sites.values()))
    # left transfer matrix
    #
    #     0      = 0     
    # i--A--j      T
    #   /\         1
    #  2  m
    #      \ 0
    #    i--A--j
    #      /
    #     2
    if len(A.size())==4:
        # assume only virtual indices are present
        a= torch.einsum('aibj->ab',A).contiguous()
    else:    
        a = torch.einsum('meifj,maibj->eafb',(A,A.conj())).contiguous().view(\
            A.size()[1]**2, A.size()[3]**2)
    a= a/a.abs().max()
    # check symmetry
    a_asymm_norm= torch.linalg.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, 2)
    # leading eigenvector is unique 
    assert torch.abs(D[0]-D[1]) > 1.0e-8, "Leading eigenvector of T not unique"
    for key,t in env.T.items():
        env.T[key]= torch.zeros(t.size(), dtype=env.dtype, device=env.device)
        env.T[key][0,0,:]= U[:,0]
//  [python]
def _orthogonalize_gram_schmidt(matrices, epsilon=0):
    """
    Apply Gram-Schmidt procedure to orthogonalize a batch of matrices.

    If epsilon is 0, this is equivalent to `torch.qr(matrices, out=(matrices, _))`,
    """
    num_cols = matrices.shape[2]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrices[:, :, i : i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input batch of matrices covers the gradients of at least one entire layer
        # in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            try:
                col /= torch.norm(col, dim=1, keepdim=True)
            except ZeroDivisionError:
                logger.error(
                    "The matrices to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 "
                    "as `orthogonalization_epsilon` in PowerSGD state."
                )
                # Recover the values from NaNs to 0s.
                col.fill_(0.0)
        else:
            col /= torch.norm(col, dim=1, keepdim=True) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrices[:, :, i + 1 :]
            rest -= torch.sum(col * rest, dim=1, keepdim=True) * col
//  [python]
def _orthogonalize_gram_schmidt(matrices, epsilon=0):
    """
    Apply Gram-Schmidt procedure to orthogonalize a batch of matrices.

    If epsilon is 0, this is equivalent to `torch.qr(matrices, out=(matrices, _))`,
    """
    num_cols = matrices.shape[2]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrices[:, :, i : i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input batch of matrices covers the gradients of at least one entire layer
        # in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            try:
                col /= torch.linalg.norm(col, dim=1, keepdim=True)
            except ZeroDivisionError:
                logger.error(
                    "The matrices to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 "
                    "as `orthogonalization_epsilon` in PowerSGD state."
                )
                # Recover the values from NaNs to 0s.
                col.fill_(0.0)
        else:
            col /= torch.linalg.norm(col, dim=1, keepdim=True) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrices[:, :, i + 1 :]
            rest -= torch.sum(col * rest, dim=1, keepdim=True) * col
//  [python]
def weighteNorm(center, neighbor):
    rel_pos = neighbor - center.unsqueeze(-2)  # (B, N, K, 3)
    rel_pos_dist = torch.norm(rel_pos, dim=-1, keepdim=True)  # (B, N, K, 1)
    rel_pos_max = torch.max(rel_pos, dim=-1, keepdim=True)[0]  # (B, N, K, 1)
    weights = torch.div((rel_pos_max - rel_pos_dist), (rel_pos_max - rel_pos_dist).sum(-2, keepdim=True))  # (B, N, K, 1)
    cov = (weights * rel_pos).transpose(-1, -2) @ rel_pos
    norm_vec = torch.symeig(cov.cpu(), eigenvectors=True).eigenvectors[..., 0]  # ascending order
    return norm_vec.to(center.device)
//  [python]
def weighteNorm(center, neighbor):
    rel_pos = neighbor - center.unsqueeze(-2)  # (B, N, K, 3)
    rel_pos_dist = torch.linalg.norm(rel_pos, dim=-1, keepdim=True)  # (B, N, K, 1)
    rel_pos_max = torch.max(rel_pos, dim=-1, keepdim=True)[0]  # (B, N, K, 1)
    weights = torch.div((rel_pos_max - rel_pos_dist), (rel_pos_max - rel_pos_dist).sum(-2, keepdim=True))  # (B, N, K, 1)
    cov = (weights * rel_pos).transpose(-1, -2) @ rel_pos
    norm_vec = torch.symeig(cov.cpu(), eigenvectors=True).eigenvectors[..., 0]  # ascending order
    return norm_vec.to(center.device)
//  [python]
def differentiable_k_means(data, k, num_iter, cluster_temp=5):
    """
    pytorch (differentiable) implementation of soft k-means clustering.
    https://github.com/bwilder0/clusternet/blob/master/models.py
    """
    # normalize x so it lies on the unit sphere
    data = torch.diag(1.0 / torch.norm(data, p=2, dim=1)) @ data
    # use kmeans++ initialization
    mu = torch.tensor(init_plus_plus(data.cpu().detach().numpy(), k), requires_grad=True).to(data)
    for _ in range(num_iter):
        # get distances between all data points and cluster centers
        dist = data @ mu.t()
        # cluster responsibilities via softmax
        r = torch.softmax(cluster_temp * dist, 1)
        # total responsibility of each cluster
        cluster_r = r.sum(dim=0)
        # mean of points in each cluster weighted by responsibility
        cluster_mean = (r.t().unsqueeze(1) @ data.expand(k, *data.shape)).squeeze(1)
        # update cluster means
        new_mu = torch.diag(1 / cluster_r) @ cluster_mean
        mu = new_mu
    dist = data @ mu.t()
    r = torch.softmax(cluster_temp * dist, 1)
    return mu, r, dist
//  [python]
def differentiable_k_means(data, k, num_iter, cluster_temp=5):
    """
    pytorch (differentiable) implementation of soft k-means clustering.
    https://github.com/bwilder0/clusternet/blob/master/models.py
    """
    # normalize x so it lies on the unit sphere
    data = torch.diag(1.0 / torch.linalg.norm(data, p=2, dim=1)) @ data
    # use kmeans++ initialization
    mu = torch.tensor(init_plus_plus(data.cpu().detach().numpy(), k), requires_grad=True).to(data)
    for _ in range(num_iter):
        # get distances between all data points and cluster centers
        dist = data @ mu.t()
        # cluster responsibilities via softmax
        r = torch.softmax(cluster_temp * dist, 1)
        # total responsibility of each cluster
        cluster_r = r.sum(dim=0)
        # mean of points in each cluster weighted by responsibility
        cluster_mean = (r.t().unsqueeze(1) @ data.expand(k, *data.shape)).squeeze(1)
        # update cluster means
        new_mu = torch.diag(1 / cluster_r) @ cluster_mean
        mu = new_mu
    dist = data @ mu.t()
    r = torch.softmax(cluster_temp * dist, 1)
    return mu, r, dist
//  [python]
def evenly_distributed_rotation(n, random_seed=None):
    """
    uniformly sample N examples on a sphere
    """
    def normalize(vector, dim: int = -1):
        return vector / torch.norm(vector, p=2.0, dim=dim, keepdim=True)

    if random_seed is not None:
        torch.manual_seed(random_seed) # fix the sampling of viewpoints for reproducing evaluation

    indices = torch.arange(0, n, dtype=torch.float32) + 0.5

    phi = torch.acos(1 - 2 * indices / n)
    theta = math.pi * (1 + 5 ** 0.5) * indices
    points = torch.stack([
        torch.cos(theta) * torch.sin(phi), 
        torch.sin(theta) * torch.sin(phi), 
        torch.cos(phi),], dim=1)
    forward = -points
    
    down = normalize(torch.randn(n, 3), dim=1)
    right = normalize(torch.cross(down, forward))
    down = normalize(torch.cross(forward, right))
    R_mat = torch.stack([right, down, forward], dim=1)
    return R_mat
//  [python]
def evenly_distributed_rotation(n, random_seed=None):
    """
    uniformly sample N examples on a sphere
    """
    def normalize(vector, dim: int = -1):
        return vector / torch.linalg.norm(vector, ord=2, dim=dim, keepdim=True)

    if random_seed is not None:
        torch.manual_seed(random_seed) # fix the sampling of viewpoints for reproducing evaluation

    indices = torch.arange(0, n, dtype=torch.float32) + 0.5

    phi = torch.acos(1 - 2 * indices / n)
    theta = math.pi * (1 + 5 ** 0.5) * indices
    points = torch.stack([
        torch.cos(theta) * torch.sin(phi), 
        torch.sin(theta) * torch.sin(phi), 
        torch.cos(phi),], dim=1)
    forward = -points
    
    down = normalize(torch.randn(n, 3), dim=1)
    right = normalize(torch.cross(down, forward))
    down = normalize(torch.cross(forward, right))
    R_mat = torch.stack([right, down, forward], dim=1)
    return R_mat
//  [python]
def create_bpr_loss(self, users, pos_items, neg_items):
        pos_scores = (users * pos_items).sum(1)
        neg_scores = (users * neg_items).sum(1)

        mf_loss = nn.LogSigmoid()(pos_scores - neg_scores).mean()
        mf_loss = -1 * mf_loss

        regularizer = (
            torch.norm(users) ** 2
            + torch.norm(pos_items) ** 2
            + torch.norm(neg_items) ** 2
        ) / 2
        emb_loss = self.lmbd * regularizer / users.shape[0]

        return mf_loss + emb_loss, mf_loss, emb_loss
//  [python]
def create_bpr_loss(self, users, pos_items, neg_items):
        pos_scores = (users * pos_items).sum(1)
        neg_scores = (users * neg_items).sum(1)

        mf_loss = nn.LogSigmoid()(pos_scores - neg_scores).mean()
        mf_loss = -1 * mf_loss

        regularizer = (
            torch.linalg.norm(users) ** 2
            + torch.linalg.norm(pos_items) ** 2
            + torch.linalg.norm(neg_items) ** 2
        ) / 2
        emb_loss = self.lmbd * regularizer / users.shape[0]

        return mf_loss + emb_loss, mf_loss, emb_loss
//  [python]
def _init_from_ipeps_pbc(a_ket, a_bra, env, verbosity=0):
    # Left-upper corner
    #
    #     i      = C--1     
    # j--A--3      0
    #   /\
    #  2  m
    #      \ i
    #    j--A--3
    #      /
    #     2
    d_ket= a_ket.size()
    d_bra= a_bra.size()
    d_kb= [d_ket[i+1]*d_bra[i+1] for i in range(4)]
    a= torch.einsum('mijef,mijab->eafb',a_ket,a_bra).contiguous().view(d_kb[2], d_kb[3])
    with torch.no_grad():
        scale= a.abs().max()
    a= a/scale

    a_asymm_norm= torch.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, a.size()[0])
    a= torch.diag(D)

    env.C[env.keyC]= torch.zeros(env.chi,env.chi, dtype=env.dtype, device=env.device)
    env.C[env.keyC][:min(env.chi,d_kb[2]),:min(env.chi,d_kb[3])]=\
       a[:min(env.chi,d_kb[2]),:min(env.chi,d_kb[3])]

    # left transfer matrix (orientation from 1->0)
    #
    #     0      = A 0     
    # i--A--3      | T--2
    #   /\         | 1
    #  2  m
    #      \ 0
    #    i--A--3
    #      /
    #     2
    a= torch.einsum('meifg,maibc->eafbgc',(a_ket,a_bra)).contiguous().view(d_kb[0], d_kb[2], d_kb[3])
    with torch.no_grad():
        scale= a.abs().max()
    a= a/scale

    a= torch.einsum('ai,abs,bj->ijs',U,a,U.conj())
    a_asymm_norm= (a-a.permute(1,0,2).conj()).norm()
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"

    env.T[env.keyT]= torch.zeros((env.chi,env.chi,d_kb[3]), dtype=env.dtype, device=env.device)
    env.T[env.keyT][:min(env.chi,d_kb[0]),:min(env.chi,d_kb[2]),:]=\
        a[:min(env.chi,d_kb[0]),:min(env.chi,d_kb[2]),:]
//  [python]
def _init_from_ipeps_pbc(a_ket, a_bra, env, verbosity=0):
    # Left-upper corner
    #
    #     i      = C--1     
    # j--A--3      0
    #   /\
    #  2  m
    #      \ i
    #    j--A--3
    #      /
    #     2
    d_ket= a_ket.size()
    d_bra= a_bra.size()
    d_kb= [d_ket[i+1]*d_bra[i+1] for i in range(4)]
    a= torch.einsum('mijef,mijab->eafb',a_ket,a_bra).contiguous().view(d_kb[2], d_kb[3])
    with torch.no_grad():
        scale= a.abs().max()
    a= a/scale

    a_asymm_norm= torch.linalg.norm(a.conj().t()-a)
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"
    D, U= truncated_eig_sym(a, a.size()[0])
    a= torch.diag(D)

    env.C[env.keyC]= torch.zeros(env.chi,env.chi, dtype=env.dtype, device=env.device)
    env.C[env.keyC][:min(env.chi,d_kb[2]),:min(env.chi,d_kb[3])]=\
       a[:min(env.chi,d_kb[2]),:min(env.chi,d_kb[3])]

    # left transfer matrix (orientation from 1->0)
    #
    #     0      = A 0     
    # i--A--3      | T--2
    #   /\         | 1
    #  2  m
    #      \ 0
    #    i--A--3
    #      /
    #     2
    a= torch.einsum('meifg,maibc->eafbgc',(a_ket,a_bra)).contiguous().view(d_kb[0], d_kb[2], d_kb[3])
    with torch.no_grad():
        scale= a.abs().max()
    a= a/scale

    a= torch.einsum('ai,abs,bj->ijs',U,a,U.conj())
    a_asymm_norm= (a-a.permute(1,0,2).conj()).norm()
    assert a_asymm_norm/a.abs().max() < 1.0e-8, "a is not symmetric"

    env.T[env.keyT]= torch.zeros((env.chi,env.chi,d_kb[3]), dtype=env.dtype, device=env.device)
    env.T[env.keyT][:min(env.chi,d_kb[0]),:min(env.chi,d_kb[2]),:]=\
        a[:min(env.chi,d_kb[0]),:min(env.chi,d_kb[2]),:]
//  [python]
def weighted_8points(x_in, logits, device):
    # x_in: batch * 1 * N * 4
    x_shp = x_in.shape
    # Turn into weights for each sample
    weights = torch.relu(torch.tanh(logits))
    x_in = x_in.squeeze(1)

    # Make input data (num_img_pair x num_corr x 4)
    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1)

    # Create the matrix to be used for the eight-point algorithm
    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1)
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1), wX)

    # Recover essential matrix from self-adjoing eigen
    v = batch_symeig(XwX, device)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def weighted_8points(x_in, logits, device):
    # x_in: batch * 1 * N * 4
    x_shp = x_in.shape
    # Turn into weights for each sample
    weights = torch.relu(torch.tanh(logits))
    x_in = x_in.squeeze(1)

    # Make input data (num_img_pair x num_corr x 4)
    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1)

    # Create the matrix to be used for the eight-point algorithm
    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1)
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1), wX)

    # Recover essential matrix from self-adjoing eigen
    v = batch_symeig(XwX, device)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.linalg.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def special_gramschmidt(M, epsilon=0):
    """
    Returns the 3x3 rotation matrix obtained by Gram-Schmidt orthonormalization of two 3D input vectors (first two columns of input matrix M).

    Args:
        M (...x3xN tensor): batch of 3xN matrices, with N >= 2.
            Only the first two columns of the matrices are used for orthonormalization.
        epsilon (float >= 0): optional clamping value to avoid returning *Not-a-Number* values in case of ill-defined input.
    Returns:
        batch of rotation matrices (...x3x3 tensor).
    Warning:
        In case of ill-defined input (colinear input column vectors), the output will not be a rotation matrix.
    """
    M, batch_shape = roma.internal.flatten_batch_dims(M, -3)
    assert(M.dim() == 3)
    x = M[:,:,0]
    y = M[:,:,1]
    x = x / torch.clamp_min(torch.norm(x, dim=-1, keepdim=True), epsilon)
    y = y - torch.sum(x*y, dim=-1, keepdim=True) * x
    y = y / torch.clamp_min(torch.norm(y, dim=-1, keepdim=True), epsilon)
    z = torch.cross(x,y, dim=-1)    
    R = torch.stack((x, y, z), dim=-1)
    return roma.internal.unflatten_batch_dims(R, batch_shape)
//  [python]
def special_gramschmidt(M, epsilon=0):
    """
    Returns the 3x3 rotation matrix obtained by Gram-Schmidt orthonormalization of two 3D input vectors (first two columns of input matrix M).

    Args:
        M (...x3xN tensor): batch of 3xN matrices, with N >= 2.
            Only the first two columns of the matrices are used for orthonormalization.
        epsilon (float >= 0): optional clamping value to avoid returning *Not-a-Number* values in case of ill-defined input.
    Returns:
        batch of rotation matrices (...x3x3 tensor).
    Warning:
        In case of ill-defined input (colinear input column vectors), the output will not be a rotation matrix.
    """
    M, batch_shape = roma.internal.flatten_batch_dims(M, -3)
    assert(M.dim() == 3)
    x = M[:,:,0]
    y = M[:,:,1]
    x = x / torch.clamp_min(torch.linalg.norm(x, dim=-1, keepdim=True), epsilon)
    y = y - torch.sum(x*y, dim=-1, keepdim=True) * x
    y = y / torch.clamp_min(torch.linalg.norm(y, dim=-1, keepdim=True), epsilon)
    z = torch.cross(x,y, dim=-1)    
    R = torch.stack((x, y, z), dim=-1)
    return roma.internal.unflatten_batch_dims(R, batch_shape)
//  [python]
def rotvec_to_unitquat(rotvec):
    """
    Converts rotation vector into unit quaternion representation.

    Args:
        rotvec (...x3 tensor): batch of rotation vectors.
    Returns:
        batch of unit quaternions (...x4 tensor, XYZW convention).
    """
    rotvec, batch_shape = roma.internal.flatten_batch_dims(rotvec, end_dim=-2)
    num_rotations, D = rotvec.shape
    assert D == 3, "Input should be a Bx3 tensor."

    # Adapted from SciPy:
    # https://github.com/scipy/scipy/blob/adc4f4f7bab120ccfab9383aba272954a0a12fb0/scipy/spatial/transform/rotation.py#L621
    
    norms = torch.norm(rotvec, dim=-1)
    small_angle = (norms <= 1e-3)
    large_angle = ~small_angle

    scale = torch.empty((num_rotations,), device=rotvec.device, dtype=rotvec.dtype)
    scale[small_angle] = (0.5 - norms[small_angle] ** 2 / 48 +
                          norms[small_angle] ** 4 / 3840)
    scale[large_angle] = (torch.sin(norms[large_angle] / 2) /
                          norms[large_angle])

    quat = torch.empty((num_rotations, 4), device=rotvec.device, dtype=rotvec.dtype)
    quat[:, :3] = scale[:, None] * rotvec
    quat[:, 3] = torch.cos(norms / 2)
    return roma.internal.unflatten_batch_dims(quat, batch_shape)
//  [python]
def rotvec_to_unitquat(rotvec):
    """
    Converts rotation vector into unit quaternion representation.

    Args:
        rotvec (...x3 tensor): batch of rotation vectors.
    Returns:
        batch of unit quaternions (...x4 tensor, XYZW convention).
    """
    rotvec, batch_shape = roma.internal.flatten_batch_dims(rotvec, end_dim=-2)
    num_rotations, D = rotvec.shape
    assert D == 3, "Input should be a Bx3 tensor."

    # Adapted from SciPy:
    # https://github.com/scipy/scipy/blob/adc4f4f7bab120ccfab9383aba272954a0a12fb0/scipy/spatial/transform/rotation.py#L621
    
    norms = torch.linalg.norm(rotvec, dim=-1)
    small_angle = (norms <= 1e-3)
    large_angle = ~small_angle

    scale = torch.empty((num_rotations,), device=rotvec.device, dtype=rotvec.dtype)
    scale[small_angle] = (0.5 - norms[small_angle] ** 2 / 48 +
                          norms[small_angle] ** 4 / 3840)
    scale[large_angle] = (torch.sin(norms[large_angle] / 2) /
                          norms[large_angle])

    quat = torch.empty((num_rotations, 4), device=rotvec.device, dtype=rotvec.dtype)
    quat[:, :3] = scale[:, None] * rotvec
    quat[:, 3] = torch.cos(norms / 2)
    return roma.internal.unflatten_batch_dims(quat, batch_shape)
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input:
            - M:           [bs, num_corr, num_corr] the compatibility matrix
            - leading_eig: [bs, num_corr]           the leading eigenvector of matrix M
        Output:
            - confidence
        """
        if method == 'eig_value':
            # max eigenvalue as the confidence (Rayleigh quotient)
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            # max eigenvalue / second max eigenvalue as the confidence
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            # compute the second largest eigen-value
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.norm(solution, dim=1, keepdim=True) + 1e-6)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = (second_eig[:, None, :] @ B @ second_eig[:, :, None]) / (
                        second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            # max xMx as the confidence (x is the binary solution)
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input:
            - M:           [bs, num_corr, num_corr] the compatibility matrix
            - leading_eig: [bs, num_corr]           the leading eigenvector of matrix M
        Output:
            - confidence
        """
        if method == 'eig_value':
            # max eigenvalue as the confidence (Rayleigh quotient)
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            # max eigenvalue / second max eigenvalue as the confidence
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            # compute the second largest eigen-value
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.linalg.norm(solution, dim=1, keepdim=True) + 1e-6)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = (second_eig[:, None, :] @ B @ second_eig[:, :, None]) / (
                        second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            # max xMx as the confidence (x is the binary solution)
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, it_num, weights=None):
        """
        Perform post refinement using the initial transformation matrix.
        https://link.springer.com/article/10.1007/s10589-014-9643-2
        Input
            - initial_trans: [bs, 4, 4]
            - src_keypts:    [bs, num_corr, 3]
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1

        if self.inlier_threshold == 0.10:  # for 3DMatch
            inlier_threshold_list = [0.10] * it_num
        else:  # for KITTI
            inlier_threshold_list = [1.2] * it_num

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)

            L2_dis = torch.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
            )

        return initial_trans
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, it_num, weights=None):
        """
        Perform post refinement using the initial transformation matrix.
        https://link.springer.com/article/10.1007/s10589-014-9643-2
        Input
            - initial_trans: [bs, 4, 4]
            - src_keypts:    [bs, num_corr, 3]
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1

        if self.inlier_threshold == 0.10:  # for 3DMatch
            inlier_threshold_list = [0.10] * it_num
        else:  # for KITTI
            inlier_threshold_list = [1.2] * it_num

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)

            L2_dis = torch.linalg.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
            )

        return initial_trans
//  [python]
def test_rope_streaming():
    set_efficient_attention_backend('torch')
    torch.manual_seed(1234)
    tr = StreamingTransformer(
        16, 4, 2, causal=True, dropout=0.,
        custom=True, positional_embedding='rope')
    tr.eval()
    steps = 12
    x = torch.randn(3, steps, 16)

    ref = tr(x)

    with tr.streaming():
        outs = []
        frame_sizes = [1] * steps

        for frame_size in frame_sizes:
            frame = x[:, :frame_size]
            x = x[:, frame_size:]
            outs.append(tr(frame))

    out = torch.cat(outs, dim=1)
    assert list(out.shape) == [3, steps, 16]
    delta = torch.norm(out - ref) / torch.norm(out)
    assert delta < 1e-6, delta
//  [python]
def test_rope_streaming():
    set_efficient_attention_backend('torch')
    torch.manual_seed(1234)
    tr = StreamingTransformer(
        16, 4, 2, causal=True, dropout=0.,
        custom=True, positional_embedding='rope')
    tr.eval()
    steps = 12
    x = torch.randn(3, steps, 16)

    ref = tr(x)

    with tr.streaming():
        outs = []
        frame_sizes = [1] * steps

        for frame_size in frame_sizes:
            frame = x[:, :frame_size]
            x = x[:, frame_size:]
            outs.append(tr(frame))

    out = torch.cat(outs, dim=1)
    assert list(out.shape) == [3, steps, 16]
    delta = torch.linalg.norm(out - ref) / torch.linalg.norm(out)
    assert delta < 1e-6, delta
//  [python]
def rotvec_to_rotmat(rotvec: torch.Tensor, epsilon=1e-6) -> torch.Tensor:
    """
    Converts rotation vector to rotation matrix representation.
    Conversion uses Rodrigues formula in general, and a first order approximation for small angles.
    
    Args:
        rotvec (...x3 tensor): batch of rotation vectors.
        epsilon (float): small angle threshold.
    Returns:
        batch of rotation matrices (...x3x3 tensor).        
    """
    rotvec, batch_shape = roma.internal.flatten_batch_dims(rotvec, end_dim=-2)
    batch_size, D = rotvec.shape
    assert(D == 3), "Input should be a Bx3 tensor."

    # Rotation angle
    theta = torch.norm(rotvec, dim=-1)
    is_angle_small = theta < epsilon
    
    # Rodrigues formula for angles that are not small.
    # Note: we use clamping to avoid non finite values for small angles
    # (torch.where produces nan gradients in such case).
    axis = rotvec / torch.clamp_min(theta[...,None], epsilon)
    kx, ky, kz = axis[:,0], axis[:,1], axis[:,2]
    sin_theta = torch.sin(theta)
    cos_theta = torch.cos(theta)
    one_minus_cos_theta = 1 - cos_theta
    xs = kx*sin_theta
    ys = ky*sin_theta
    zs = kz*sin_theta
    xyc = kx*ky*one_minus_cos_theta
    xzc = kx*kz*one_minus_cos_theta
    yzc = ky*kz*one_minus_cos_theta
    xxc = kx**2*one_minus_cos_theta
    yyc = ky**2*one_minus_cos_theta
    zzc = kz**2*one_minus_cos_theta
    R_rodrigues = torch.stack([1 - yyc - zzc, xyc - zs, xzc + ys,
                     xyc + zs, 1 - xxc - zzc, -xs + yzc,
                     xzc - ys, xs + yzc, 1 -xxc - yyc], dim=-1).reshape(-1, 3, 3)

    # For small angles, use a first order approximation
    xs, ys, zs = rotvec[:,0], rotvec[:,1], rotvec[:,2]
    one = torch.ones_like(xs)
    R_first_order = torch.stack([one, -zs, ys,
                                 zs, one, -xs,
                                 -ys, xs, one], dim=-1).reshape(-1, 3, 3)
    # Select the appropriate expression
    R = torch.where(is_angle_small[:,None,None], R_first_order, R_rodrigues)
    return roma.internal.unflatten_batch_dims(R, batch_shape)
//  [python]
def rotvec_to_rotmat(rotvec: torch.Tensor, epsilon=1e-6) -> torch.Tensor:
    """
    Converts rotation vector to rotation matrix representation.
    Conversion uses Rodrigues formula in general, and a first order approximation for small angles.
    
    Args:
        rotvec (...x3 tensor): batch of rotation vectors.
        epsilon (float): small angle threshold.
    Returns:
        batch of rotation matrices (...x3x3 tensor).        
    """
    rotvec, batch_shape = roma.internal.flatten_batch_dims(rotvec, end_dim=-2)
    batch_size, D = rotvec.shape
    assert(D == 3), "Input should be a Bx3 tensor."

    # Rotation angle
    theta = torch.linalg.norm(rotvec, dim=-1)
    is_angle_small = theta < epsilon
    
    # Rodrigues formula for angles that are not small.
    # Note: we use clamping to avoid non finite values for small angles
    # (torch.where produces nan gradients in such case).
    axis = rotvec / torch.clamp_min(theta[...,None], epsilon)
    kx, ky, kz = axis[:,0], axis[:,1], axis[:,2]
    sin_theta = torch.sin(theta)
    cos_theta = torch.cos(theta)
    one_minus_cos_theta = 1 - cos_theta
    xs = kx*sin_theta
    ys = ky*sin_theta
    zs = kz*sin_theta
    xyc = kx*ky*one_minus_cos_theta
    xzc = kx*kz*one_minus_cos_theta
    yzc = ky*kz*one_minus_cos_theta
    xxc = kx**2*one_minus_cos_theta
    yyc = ky**2*one_minus_cos_theta
    zzc = kz**2*one_minus_cos_theta
    R_rodrigues = torch.stack([1 - yyc - zzc, xyc - zs, xzc + ys,
                     xyc + zs, 1 - xxc - zzc, -xs + yzc,
                     xzc - ys, xs + yzc, 1 -xxc - yyc], dim=-1).reshape(-1, 3, 3)

    # For small angles, use a first order approximation
    xs, ys, zs = rotvec[:,0], rotvec[:,1], rotvec[:,2]
    one = torch.ones_like(xs)
    R_first_order = torch.stack([one, -zs, ys,
                                 zs, one, -xs,
                                 -ys, xs, one], dim=-1).reshape(-1, 3, 3)
    # Select the appropriate expression
    R = torch.where(is_angle_small[:,None,None], R_first_order, R_rodrigues)
    return roma.internal.unflatten_batch_dims(R, batch_shape)
//  [python]
def newton_raphson_chunk(
    fn, x0, solver, rtol=1e-6, atol=1e-10, miter=100, throw_on_fail=False
):
    """
    Solve a nonlinear system with Newton's method with a tensor for a
    BackwardEuler type chunking operator context manager.

    Args:
      fn (function):        function that returns R, J, and the solver context
      x0 (torch.tensor):    starting point
      solver (ChunkTimeOperatorSolverContext): solver context

    Keyword Args:
      rtol (float):         nonlinear relative tolerance
      atol (float):         nonlinear absolute tolerance
      miter (int):          maximum number of nonlinear iterations
      throw_on_fail (bool): throw exception if solve fails

    Returns:
      torch.tensor:         solution to system of equations
    """
    x = x0
    R, J = fn(x)

    nR = torch.norm(R, dim=-1)
    nR0 = nR
    i = 0

    while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
        x -= solver.solve(J, R)
        R, J = fn(x)
        nR = torch.norm(R, dim=-1)
        i += 1

    if i == miter:
        if throw_on_fail:
            raise RuntimeError("Implicit solve did not succeed.")
        warnings.warn("Implicit solve did not succeed.  Results may be inaccurate...")

    return x
//  [python]
def newton_raphson_chunk(
    fn, x0, solver, rtol=1e-6, atol=1e-10, miter=100, throw_on_fail=False
):
    """
    Solve a nonlinear system with Newton's method with a tensor for a
    BackwardEuler type chunking operator context manager.

    Args:
      fn (function):        function that returns R, J, and the solver context
      x0 (torch.tensor):    starting point
      solver (ChunkTimeOperatorSolverContext): solver context

    Keyword Args:
      rtol (float):         nonlinear relative tolerance
      atol (float):         nonlinear absolute tolerance
      miter (int):          maximum number of nonlinear iterations
      throw_on_fail (bool): throw exception if solve fails

    Returns:
      torch.tensor:         solution to system of equations
    """
    x = x0
    R, J = fn(x)

    nR = torch.linalg.norm(R, dim=-1)
    nR0 = nR
    i = 0

    while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
        x -= solver.solve(J, R)
        R, J = fn(x)
        nR = torch.linalg.norm(R, dim=-1)
        i += 1

    if i == miter:
        if throw_on_fail:
            raise RuntimeError("Implicit solve did not succeed.")
        warnings.warn("Implicit solve did not succeed.  Results may be inaccurate...")

    return x
//  [python]
def soft_corr_generation(self, src_keypts, tgt_keypts, src_features, tgt_features):
        """
        Generate soft correspondences by feature matching
        Input:
            - src_keypts:      [bs, num_point, 3]  src keypoints
            - tgt_keypts:      [bs, num_point, 3]  tgt keypoints
            - src_features:    [bs, num_point, feat_dim]  src features of src keypoints
            - tgt_features:    [bs, num_point, feat_dim]  tgt features of tgt keypoints
        Output:
            - src_keypts_corr: [bs, num_corr,3 ] tgt correspondences of src keypoints
            - tgt_keypts_corr: [bs, num_corr,3 ] tgt correspondences of tgt keypoints
        """
        # normalization
        src_features = src_features / (torch.norm(src_features, dim=1, keepdim=True) + 1e-6)
        tgt_features = tgt_features / (torch.norm(tgt_features, dim=1, keepdim=True) + 1e-6)
        N_src = src_features.shape[1]
        N_tgt = tgt_features.shape[1]
        # use all point or sample points.
        if self.num_node == 'all':
            src_sel_ind = np.arange(N_src)
            tgt_sel_ind = np.arange(N_tgt)
        else:
            src_sel_ind = np.random.choice(N_src, self.num_node)
            tgt_sel_ind = np.random.choice(N_tgt, self.num_node)
        src_desc = src_features[:, src_sel_ind, :]
        tgt_desc = tgt_features[:, tgt_sel_ind, :]
        src_keypts = src_keypts[:, src_sel_ind, :]
        tgt_keypts = tgt_keypts[:, tgt_sel_ind, :]

        # match points in feature space.
        distance = torch.sqrt(2 - 2 * (src_desc[0] @ tgt_desc[0].T) + 1e-6)
        distance = distance.unsqueeze(0)
        source_idx = torch.argmin(distance[0], dim=1)
        corr = torch.cat([torch.arange(source_idx.shape[0])[:, None].cuda(), source_idx[:, None]], dim=-1)

        # generate correspondences
        src_keypts_corr = src_keypts[:, corr[:, 0]]
        tgt_keypts_corr = tgt_keypts[:, corr[:, 1]]

        return src_keypts_corr, tgt_keypts_corr
//  [python]
def soft_corr_generation(self, src_keypts, tgt_keypts, src_features, tgt_features):
        """
        Generate soft correspondences by feature matching
        Input:
            - src_keypts:      [bs, num_point, 3]  src keypoints
            - tgt_keypts:      [bs, num_point, 3]  tgt keypoints
            - src_features:    [bs, num_point, feat_dim]  src features of src keypoints
            - tgt_features:    [bs, num_point, feat_dim]  tgt features of tgt keypoints
        Output:
            - src_keypts_corr: [bs, num_corr,3 ] tgt correspondences of src keypoints
            - tgt_keypts_corr: [bs, num_corr,3 ] tgt correspondences of tgt keypoints
        """
        # normalization
        src_features = src_features / (torch.linalg.norm(src_features, dim=1, keepdim=True) + 1e-6)
        tgt_features = tgt_features / (torch.linalg.norm(tgt_features, dim=1, keepdim=True) + 1e-6)
        N_src = src_features.shape[1]
        N_tgt = tgt_features.shape[1]
        # use all point or sample points.
        if self.num_node == 'all':
            src_sel_ind = np.arange(N_src)
            tgt_sel_ind = np.arange(N_tgt)
        else:
            src_sel_ind = np.random.choice(N_src, self.num_node)
            tgt_sel_ind = np.random.choice(N_tgt, self.num_node)
        src_desc = src_features[:, src_sel_ind, :]
        tgt_desc = tgt_features[:, tgt_sel_ind, :]
        src_keypts = src_keypts[:, src_sel_ind, :]
        tgt_keypts = tgt_keypts[:, tgt_sel_ind, :]

        # match points in feature space.
        distance = torch.sqrt(2 - 2 * (src_desc[0] @ tgt_desc[0].T) + 1e-6)
        distance = distance.unsqueeze(0)
        source_idx = torch.argmin(distance[0], dim=1)
        corr = torch.cat([torch.arange(source_idx.shape[0])[:, None].cuda(), source_idx[:, None]], dim=-1)

        # generate correspondences
        src_keypts_corr = src_keypts[:, corr[:, 0]]
        tgt_keypts_corr = tgt_keypts[:, corr[:, 1]]

        return src_keypts_corr, tgt_keypts_corr
//  [python]
def test_rope_streaming_past_context():
    set_efficient_attention_backend('torch')
    torch.manual_seed(1234)

    for context in [None, 10]:
        tr = StreamingTransformer(
            16, 4, 1 if context else 2,
            causal=True, past_context=context, custom=True,
            dropout=0., positional_embedding='rope')
        tr.eval()

        steps = 20
        x = torch.randn(3, steps, 16)
        ref = tr(x)

        with tr.streaming():
            outs = []
            frame_sizes = [1] * steps

            for frame_size in frame_sizes:
                frame = x[:, :frame_size]
                x = x[:, frame_size:]
                outs.append(tr(frame))

        out = torch.cat(outs, dim=1)
        assert list(out.shape) == [3, steps, 16]
        delta = torch.norm(out - ref) / torch.norm(out)
        assert delta < 1e-6, delta
//  [python]
def test_rope_streaming_past_context():
    set_efficient_attention_backend('torch')
    torch.manual_seed(1234)

    for context in [None, 10]:
        tr = StreamingTransformer(
            16, 4, 1 if context else 2,
            causal=True, past_context=context, custom=True,
            dropout=0., positional_embedding='rope')
        tr.eval()

        steps = 20
        x = torch.randn(3, steps, 16)
        ref = tr(x)

        with tr.streaming():
            outs = []
            frame_sizes = [1] * steps

            for frame_size in frame_sizes:
                frame = x[:, :frame_size]
                x = x[:, frame_size:]
                outs.append(tr(frame))

        out = torch.cat(outs, dim=1)
        assert list(out.shape) == [3, steps, 16]
        delta = torch.linalg.norm(out - ref) / torch.linalg.norm(out)
        assert delta < 1e-6, delta
//  [python]
def capsule_volume_batch(c1: torch.Tensor, c2: torch.Tensor, r: torch.Tensor):
    """Compute the volume of a capsule given its end points and radius.

    Args:
        c1 [batch x 3]: First end point.
        c2 [batch x 3]: Second end point.
        r [batch]: Radius of the capsule.

    Returns:
        [batch]: Volume of the capsule.
    """

    h = torch.norm(c2 - c1, dim=1)
    return np.pi * h * (r**2) + (4 / 3) * np.pi * (r**3)
//  [python]
def capsule_volume_batch(c1: torch.Tensor, c2: torch.Tensor, r: torch.Tensor):
    """Compute the volume of a capsule given its end points and radius.

    Args:
        c1 [batch x 3]: First end point.
        c2 [batch x 3]: Second end point.
        r [batch]: Radius of the capsule.

    Returns:
        [batch]: Volume of the capsule.
    """

    h = torch.linalg.norm(c2 - c1, dim=1)
    return np.pi * h * (r**2) + (4 / 3) * np.pi * (r**3)
//  [python]
def hierarchical_corr_filtering(self, src_keypts, tgt_keypts):
        """
        Estimate the transformation by SVD and hierarchical correspondence filtering.
        Input:
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
        Output:
            - pred_trans:   [bs, 4, 4], the predicted transformation matrix.
            - pred_labels:  [bs, num_corr], the predicted inlier/outlier label (0,1)
        """
        bs, num_corr = src_keypts.shape[0], tgt_keypts.shape[1]

        #################################
        # downsample points
        #################################
        if num_corr > self.max_points:
            src_keypts = src_keypts[:, :self.max_points, :]
            tgt_keypts = tgt_keypts[:, :self.max_points, :]
            num_corr = self.max_points

        #################################
        # compute cross dist
        #################################
        src_dist = torch.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
        target_dist = torch.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]), dim=-1)
        cross_dist = torch.abs(src_dist - target_dist)

        #################################
        # compute first order measure
        #################################
        SC_dist_thre = self.d_thre
        SC_measure = torch.clamp(1.0 - cross_dist ** 2 / SC_dist_thre ** 2, min=0)
        hard_SC_measure = (cross_dist < SC_dist_thre).float()

        #################################
        # select reliable seed correspondences
        #################################
        confidence = self.cal_leading_eigenvector(SC_measure, method='power')
        seeds = self.pick_seeds(src_dist, confidence, R=self.nms_radius, max_num=int(num_corr * self.ratio))

        #################################
        # compute second order measure
        #################################
        SC2_dist_thre = self.d_thre / 2
        hard_SC_measure_tight = (cross_dist < SC2_dist_thre).float()
        seed_hard_SC_measure = hard_SC_measure.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_corr))
        seed_hard_SC_measure_tight = hard_SC_measure_tight.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_corr))

        SC2_measure = torch.matmul(seed_hard_SC_measure_tight, hard_SC_measure_tight) * seed_hard_SC_measure

        #################################
        # compute the seed-wise transformations and select the best one
        #################################
        final_trans = self.trans_estimator(SC2_measure, src_keypts, tgt_keypts)

        #################################
        # refine the result by recomputing the transformation over the whole set
        #################################
        final_trans = self.post_refinement(final_trans, src_keypts, tgt_keypts, 20)

        return final_trans