jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
        target_dist = torch.linalg.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
//  [python]
def point_capsule_distance_batch(p: torch.Tensor, c1: torch.Tensor, c2: torch.Tensor, r: torch.Tensor):
    """Compute the distance between a point and a capsule given its end points and radius.

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

    Returns:
        [batch]: Distance between the point and the capsule.
    """

    pc1 = p - c1
    c2c1 = c2 - c1
    h = torch.clamp((pc1 * c2c1).sum(dim=1) / (c2c1 * c2c1).sum(dim=1), 0, 1)
    return torch.norm(pc1 - h.unsqueeze(1) * c2c1, dim=1) - r
//  [python]
def point_capsule_distance_batch(p: torch.Tensor, c1: torch.Tensor, c2: torch.Tensor, r: torch.Tensor):
    """Compute the distance between a point and a capsule given its end points and radius.

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

    Returns:
        [batch]: Distance between the point and the capsule.
    """

    pc1 = p - c1
    c2c1 = c2 - c1
    h = torch.clamp((pc1 * c2c1).sum(dim=1) / (c2c1 * c2c1).sum(dim=1), 0, 1)
    return torch.linalg.norm(pc1 - h.unsqueeze(1) * c2c1, dim=1) - r
//  [python]
def cosine_similarity(test: Tensor, train: Tensor, replace_nan: int = 0) -> Tensor:
    r"""
    Calculates the pairwise cosine similarity for batches of feature vectors.
    Tensors test and train have shape (batch_size_1, *), and (batch_size_2, *).
    Returns pairwise cosine similarity Tensor of shape (batch_size_1, batch_size_2).
    """
    test = test.view(test.shape[0], -1)
    train = train.view(train.shape[0], -1)

    if common._parse_version(torch.__version__) <= (1, 6, 0):
        test_norm = torch.norm(test, p=None, dim=1, keepdim=True)
        train_norm = torch.norm(train, p=None, dim=1, keepdim=True)
    else:
        test_norm = torch.linalg.norm(test, ord=2, dim=1, keepdim=True)
        train_norm = torch.linalg.norm(train, ord=2, dim=1, keepdim=True)

    test = torch.where(test_norm != 0.0, test / test_norm, Tensor([replace_nan]))
    train = torch.where(train_norm != 0.0, train / train_norm, Tensor([replace_nan])).T

    similarity = torch.mm(test, train)
    return similarity
//  [python]
def cosine_similarity(test: Tensor, train: Tensor, replace_nan: int = 0) -> Tensor:
    r"""
    Calculates the pairwise cosine similarity for batches of feature vectors.
    Tensors test and train have shape (batch_size_1, *), and (batch_size_2, *).
    Returns pairwise cosine similarity Tensor of shape (batch_size_1, batch_size_2).
    """
    test = test.view(test.shape[0], -1)
    train = train.view(train.shape[0], -1)

    test_norm = torch.linalg.norm(test, ord=2, dim=1, keepdim=True)
    train_norm = torch.linalg.norm(train, ord=2, dim=1, keepdim=True)

    test = torch.where(test_norm != 0.0, test / test_norm, Tensor([replace_nan]))
    train = torch.where(train_norm != 0.0, train / train_norm, Tensor([replace_nan])).T

    similarity = torch.mm(test, train)
    return similarity
//  [python]
def test_percentile_clipping(gtype):
    gnorm_vec1 = torch.zeros(100, device="cuda")
    gnorm_vec2 = torch.zeros(100, device="cuda")
    n = 4
    step = 0
    percentile = 5
    for i in range(k):
        step += 1
        g = torch.randn(n, n, dtype=gtype, device="cuda")
        gnorm1, clip2, gnorm_scale = F.percentile_clipping(g, gnorm_vec2, step, percentile=percentile)
        assert gnorm_scale == 1.0 if gnorm1 < clip2 else clip2 / gnorm1

        gnorm2 = torch.norm(g.float())
        if step == 1:
            gnorm_vec1[:] = gnorm2
        else:
            gnorm_vec1[step % 100] = gnorm2

        vals, idx = torch.sort(gnorm_vec1)
        clip1 = vals[percentile]

        torch.testing.assert_close(gnorm_vec1, torch.sqrt(gnorm_vec2))
        torch.testing.assert_close(clip1, clip2)
        torch.testing.assert_close(gnorm1, gnorm2)
//  [python]
def test_percentile_clipping(gtype):
    gnorm_vec1 = torch.zeros(100, device="cuda")
    gnorm_vec2 = torch.zeros(100, device="cuda")
    n = 4
    step = 0
    percentile = 5
    for i in range(k):
        step += 1
        g = torch.randn(n, n, dtype=gtype, device="cuda")
        gnorm1, clip2, gnorm_scale = F.percentile_clipping(g, gnorm_vec2, step, percentile=percentile)
        assert gnorm_scale == 1.0 if gnorm1 < clip2 else clip2 / gnorm1

        gnorm2 = torch.linalg.norm(g.float())
        if step == 1:
            gnorm_vec1[:] = gnorm2
        else:
            gnorm_vec1[step % 100] = gnorm2

        vals, idx = torch.sort(gnorm_vec1)
        clip1 = vals[percentile]

        torch.testing.assert_close(gnorm_vec1, torch.sqrt(gnorm_vec2))
        torch.testing.assert_close(clip1, clip2)
        torch.testing.assert_close(gnorm1, gnorm2)
//  [python]
def dist_mat(self, x, y=None):

        try:
            x = th.from_numpy(x)
        except TypeError:
            x = x

        dist = th.norm(x[:, None] - x, dim=2, p=2)
        return dist / dist.max()
//  [python]
def dist_mat(self, x, y=None):

        try:
            x = th.from_numpy(x)
        except TypeError:
            x = x

        dist = th.linalg.norm(x[:, None] - x, dim=2, ord=2)
        return dist / dist.max()
//  [python]
def backward(ctx, grad_output):
        grad_Z = grad_output
        F, S, Q_F, Q_S, Z, G, X, gamma = ctx.saved_tensors
        FF = F.t() @ F
        FF_norm = torch.norm(FF, p='fro')
        R = G * (Q_F.t() @ grad_Z @ Q_S)
        R = Q_F @ R @ Q_S.t() @ torch.sparse.mm(S, Z.t())
        scalar_1 = gamma * (1/(FF_norm+epsilon_F))
        scalar_2 = torch.sum(FF * R)
        scalar_2 = 2 * scalar_2 * (1/(FF_norm**2 + epsilon_F * FF_norm))
        grad_F = (R + R.t()) - scalar_2 * FF
        grad_F = scalar_1 * (F @ grad_F)
        grad_X = None
        return grad_X, grad_F, None, None, None, None
//  [python]
def backward(ctx, grad_output):
        grad_Z = grad_output
        F, S, Q_F, Q_S, Z, G, X, gamma = ctx.saved_tensors
        FF = F.t() @ F
        FF_norm = torch.linalg.norm(FF, ord='fro')
        R = G * (Q_F.t() @ grad_Z @ Q_S)
        R = Q_F @ R @ Q_S.t() @ torch.sparse.mm(S, Z.t())
        scalar_1 = gamma * (1/(FF_norm+epsilon_F))
        scalar_2 = torch.sum(FF * R)
        scalar_2 = 2 * scalar_2 * (1/(FF_norm**2 + epsilon_F * FF_norm))
        grad_F = (R + R.t()) - scalar_2 * FF
        grad_F = scalar_1 * (F @ grad_F)
        grad_X = None
        return grad_X, grad_F, None, None, None, None
//  [python]
def euler2rot(euler):
    """Converts a [B,3] torch.tensor euler angle vector to a [B,4,4] torch.tensor pose"""

    euler_norm = torch.norm(euler, 2, 2, True)
    axis = euler / (euler_norm + 1e-7)

    cos_a = torch.cos(euler_norm)
    sin_a = torch.sin(euler_norm)
    cos1_a = 1 - cos_a

    x = axis[..., 0].unsqueeze(1)
    y = axis[..., 1].unsqueeze(1)
    z = axis[..., 2].unsqueeze(1)

    x_sin = x * sin_a
    y_sin = y * sin_a
    z_sin = z * sin_a
    x_cos1 = x * cos1_a
    y_cos1 = y * cos1_a
    z_cos1 = z * cos1_a

    xx_cos1 = x * x_cos1
    yy_cos1 = y * y_cos1
    zz_cos1 = z * z_cos1
    xy_cos1 = x * y_cos1
    yz_cos1 = y * z_cos1
    zx_cos1 = z * x_cos1

    batch_size = euler.shape[0]
    rot = torch.zeros((batch_size, 4, 4)).to(device=euler.device)

    rot[:, 0, 0] = torch.squeeze(xx_cos1 + cos_a)
    rot[:, 0, 1] = torch.squeeze(xy_cos1 - z_sin)
    rot[:, 0, 2] = torch.squeeze(zx_cos1 + y_sin)
    rot[:, 1, 0] = torch.squeeze(xy_cos1 + z_sin)
    rot[:, 1, 1] = torch.squeeze(yy_cos1 + cos_a)
    rot[:, 1, 2] = torch.squeeze(yz_cos1 - x_sin)
    rot[:, 2, 0] = torch.squeeze(zx_cos1 - y_sin)
    rot[:, 2, 1] = torch.squeeze(yz_cos1 + x_sin)
    rot[:, 2, 2] = torch.squeeze(zz_cos1 + cos_a)
    rot[:, 3, 3] = 1

    return rot
//  [python]
def euler2rot(euler):
    """Converts a [B,3] torch.tensor euler angle vector to a [B,4,4] torch.tensor pose"""

    euler_norm = torch.linalg.norm(euler, 2, 2, True)
    axis = euler / (euler_norm + 1e-7)

    cos_a = torch.cos(euler_norm)
    sin_a = torch.sin(euler_norm)
    cos1_a = 1 - cos_a

    x = axis[..., 0].unsqueeze(1)
    y = axis[..., 1].unsqueeze(1)
    z = axis[..., 2].unsqueeze(1)

    x_sin = x * sin_a
    y_sin = y * sin_a
    z_sin = z * sin_a
    x_cos1 = x * cos1_a
    y_cos1 = y * cos1_a
    z_cos1 = z * cos1_a

    xx_cos1 = x * x_cos1
    yy_cos1 = y * y_cos1
    zz_cos1 = z * z_cos1
    xy_cos1 = x * y_cos1
    yz_cos1 = y * z_cos1
    zx_cos1 = z * x_cos1

    batch_size = euler.shape[0]
    rot = torch.zeros((batch_size, 4, 4)).to(device=euler.device)

    rot[:, 0, 0] = torch.squeeze(xx_cos1 + cos_a)
    rot[:, 0, 1] = torch.squeeze(xy_cos1 - z_sin)
    rot[:, 0, 2] = torch.squeeze(zx_cos1 + y_sin)
    rot[:, 1, 0] = torch.squeeze(xy_cos1 + z_sin)
    rot[:, 1, 1] = torch.squeeze(yy_cos1 + cos_a)
    rot[:, 1, 2] = torch.squeeze(yz_cos1 - x_sin)
    rot[:, 2, 0] = torch.squeeze(zx_cos1 - y_sin)
    rot[:, 2, 1] = torch.squeeze(yz_cos1 + x_sin)
    rot[:, 2, 2] = torch.squeeze(zz_cos1 + cos_a)
    rot[:, 3, 3] = 1

    return rot
//  [python]
def _solve_pcg(
        self,
        hvp: Callable[[torch.Tensor], torch.Tensor],
        b: torch.Tensor,
        *,
        tol: float,
        x0: Optional[torch.Tensor] = None,
        maxiter: Optional[int] = None,
    ) -> torch.Tensor:
        r"""
        Conjugate gradient solver for the Hessian vector product.

        Args:
            hvp: A callable Hvp, operating with tensors of size N.
            b: A vector or matrix, the right hand side of the equation \(Hx = b\).
            x0: Initial guess for hvp.
            rtol: Maximum relative tolerance of result.
            atol: Absolute tolerance of result.
            maxiter: Maximum number of iterations. If None, defaults to 10*len(b).

        Returns:
            A tensor with the solution of \(Ax=b\).
        """

        if x0 is None:
            x0 = torch.clone(b)
        if maxiter is None:
            maxiter = len(b) * 10

        x = x0

        r0 = b - hvp(x)

        if self.pre_conditioner is not None:
            p = z0 = self.pre_conditioner.solve(r0)
        else:
            p = z0 = r0

        for k in range(maxiter):
            if torch.norm(r0) < tol:
                break
            Ap = hvp(p)
            alpha = torch.dot(r0, z0) / torch.dot(p, Ap)
            x += alpha * p
            r = r0 - alpha * Ap

            if self.pre_conditioner is not None:
                z = self.pre_conditioner.solve(r)
            else:
                z = r

            beta = torch.dot(r, z) / torch.dot(r0, z0)

            r0 = r
            p = z + beta * p
            z0 = z

        return x
//  [python]
def _solve_pcg(
        self,
        hvp: Callable[[torch.Tensor], torch.Tensor],
        b: torch.Tensor,
        *,
        tol: float,
        x0: Optional[torch.Tensor] = None,
        maxiter: Optional[int] = None,
    ) -> torch.Tensor:
        r"""
        Conjugate gradient solver for the Hessian vector product.

        Args:
            hvp: A callable Hvp, operating with tensors of size N.
            b: A vector or matrix, the right hand side of the equation \(Hx = b\).
            x0: Initial guess for hvp.
            rtol: Maximum relative tolerance of result.
            atol: Absolute tolerance of result.
            maxiter: Maximum number of iterations. If None, defaults to 10*len(b).

        Returns:
            A tensor with the solution of \(Ax=b\).
        """

        if x0 is None:
            x0 = torch.clone(b)
        if maxiter is None:
            maxiter = len(b) * 10

        x = x0

        r0 = b - hvp(x)

        if self.pre_conditioner is not None:
            p = z0 = self.pre_conditioner.solve(r0)
        else:
            p = z0 = r0

        for k in range(maxiter):
            if torch.linalg.norm(r0) < tol:
                break
            Ap = hvp(p)
            alpha = torch.dot(r0, z0) / torch.dot(p, Ap)
            x += alpha * p
            r = r0 - alpha * Ap

            if self.pre_conditioner is not None:
                z = self.pre_conditioner.solve(r)
            else:
                z = r

            beta = torch.dot(r, z) / torch.dot(r0, z0)

            r0 = r
            p = z + beta * p
            z0 = z

        return x
//  [python]
def forward(self, *input):
        src_embedding = input[0]
        tgt_embedding = input[1]
        embedding = torch.cat((src_embedding, tgt_embedding), dim=1)
        embedding = self.nn(embedding.max(dim=-1)[0])
        rotation = self.proj_rot(embedding)
        rotation = rotation / torch.norm(rotation, p=2, dim=1, keepdim=True)
        translation = self.proj_trans(embedding)
        return quat2mat(rotation), translation
//  [python]
def forward(self, *input):
        src_embedding = input[0]
        tgt_embedding = input[1]
        embedding = torch.cat((src_embedding, tgt_embedding), dim=1)
        embedding = self.nn(embedding.max(dim=-1)[0])
        rotation = self.proj_rot(embedding)
        rotation = rotation / torch.linalg.norm(rotation, p=2, dim=1, keepdim=True)
        translation = self.proj_trans(embedding)
        return quat2mat(rotation), translation
//  [python]
def function_and_gradient_full_acc_torch(A, B, V, VVT, regu, regu_norm):
    """
    Computes function value and gradient for the optimization problem with the added
    advantage of precomputation of V.dot(V.T)

    Parameters
    ----------
    A : numpy array
        Gene expression matrix.
    B : numpy array
        Filtering matrix.
    V : numpy array
        Spectral matrix.
    VVT : numpy array
        The precomputed value of V.dot(V.T)
    regu : float
        Regularization coefficient.
    regu_norm : str, optional
        Type of regularization norm ('L1' or 'L2').

    Returns
    -------
    tuple
        A tuple containing the projection over theoretic spectrum and the gradient according to 'B'.
    """
    T_0 = (A * B)
    t_1 = torch.norm(B, 1)
    functionValue = (torch.trace((((V.T).mm(T_0)).mm(T_0.T)).mm(V)) - (regu * t_1))
    gradient = ((2 * ((VVT).mm(T_0) * A)) - ((regu) * torch.sign(B)))
    return functionValue, gradient
//  [python]
def function_and_gradient_full_acc_torch(A, B, V, VVT, regu, regu_norm):
    """
    Computes function value and gradient for the optimization problem with the added
    advantage of precomputation of V.dot(V.T)

    Parameters
    ----------
    A : numpy array
        Gene expression matrix.
    B : numpy array
        Filtering matrix.
    V : numpy array
        Spectral matrix.
    VVT : numpy array
        The precomputed value of V.dot(V.T)
    regu : float
        Regularization coefficient.
    regu_norm : str, optional
        Type of regularization norm ('L1' or 'L2').

    Returns
    -------
    tuple
        A tuple containing the projection over theoretic spectrum and the gradient according to 'B'.
    """
    T_0 = (A * B)
    t_1 = torch.linalg.norm(B, 1)
    functionValue = (torch.trace((((V.T).mm(T_0)).mm(T_0.T)).mm(V)) - (regu * t_1))
    gradient = ((2 * ((VVT).mm(T_0) * A)) - ((regu) * torch.sign(B)))
    return functionValue, gradient
//  [python]
def forward(self, *input):
        src = input[0]
        tgt = input[1]
        src_embedding = input[2]
        tgt_embedding = input[3]
        batch_size, num_dims, num_points = src_embedding.size()
        src_norm = torch.norm(src_embedding, dim=1, keepdim=True)
        tgt_norm = torch.norm(tgt_embedding, dim=1, keepdim=True)
        src_topk_idx = torch.topk(src_norm, k=self.num_keypoints, dim=2, sorted=False)[1]
        tgt_topk_idx = torch.topk(tgt_norm, k=self.num_keypoints, dim=2, sorted=False)[1]
        src_keypoints_idx = src_topk_idx.repeat(1, 3, 1)
        tgt_keypoints_idx = tgt_topk_idx.repeat(1, 3, 1)
        src_embedding_idx = src_topk_idx.repeat(1, num_dims, 1)
        tgt_embedding_idx = tgt_topk_idx.repeat(1, num_dims, 1)

        src_keypoints = torch.gather(src, dim=2, index=src_keypoints_idx)
        tgt_keypoints = torch.gather(tgt, dim=2, index=tgt_keypoints_idx)
        
        src_embedding = torch.gather(src_embedding, dim=2, index=src_embedding_idx)
        tgt_embedding = torch.gather(tgt_embedding, dim=2, index=tgt_embedding_idx)
        return src_keypoints, tgt_keypoints, src_embedding, tgt_embedding
//  [python]
def forward(self, *input):
        src = input[0]
        tgt = input[1]
        src_embedding = input[2]
        tgt_embedding = input[3]
        batch_size, num_dims, num_points = src_embedding.size()
        src_norm = torch.linalg.norm(src_embedding, dim=1, keepdim=True)
        tgt_norm = torch.linalg.norm(tgt_embedding, dim=1, keepdim=True)
        src_topk_idx = torch.topk(src_norm, k=self.num_keypoints, dim=2, sorted=False)[1]
        tgt_topk_idx = torch.topk(tgt_norm, k=self.num_keypoints, dim=2, sorted=False)[1]
        src_keypoints_idx = src_topk_idx.repeat(1, 3, 1)
        tgt_keypoints_idx = tgt_topk_idx.repeat(1, 3, 1)
        src_embedding_idx = src_topk_idx.repeat(1, num_dims, 1)
        tgt_embedding_idx = tgt_topk_idx.repeat(1, num_dims, 1)

        src_keypoints = torch.gather(src, dim=2, index=src_keypoints_idx)
        tgt_keypoints = torch.gather(tgt, dim=2, index=tgt_keypoints_idx)
        
        src_embedding = torch.gather(src_embedding, dim=2, index=src_embedding_idx)
        tgt_embedding = torch.gather(tgt_embedding, dim=2, index=tgt_embedding_idx)
        return src_keypoints, tgt_keypoints, src_embedding, tgt_embedding
//  [python]
def BBS_torch(E, prev_E, I, ones_m, iterNum=1000):
    ''' Bregmanian Bi-Stochastication algorithm as described inL
    https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/KAIS_BBS_final.pdf
    :param E: permutation matrix
    :param iterNum:iteration number
    :return:Bi-Stochastic matrix
    '''
    n = E.shape[0]
    for i in range(iterNum):
        if i % 10 == 1:
            prev_E = torch.clone(E)
        ones_E = ones_m.mm(E)
        E = E + (1 / n) * (I - E + (1 / n) * (ones_E)).mm(ones_m) - (1 / n) * ones_E
        E = torch.clip(E, min=0, max=1)  # E,E.shape[0],E.shape[0],0)
        if i % 10 == 1:
            if torch.norm(E - prev_E) < ((10e-7) * n * n):
                break
    return E
//  [python]
def BBS_torch(E, prev_E, I, ones_m, iterNum=1000):
    ''' Bregmanian Bi-Stochastication algorithm as described inL
    https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/KAIS_BBS_final.pdf
    :param E: permutation matrix
    :param iterNum:iteration number
    :return:Bi-Stochastic matrix
    '''
    n = E.shape[0]
    for i in range(iterNum):
        if i % 10 == 1:
            prev_E = torch.clone(E)
        ones_E = ones_m.mm(E)
        E = E + (1 / n) * (I - E + (1 / n) * (ones_E)).mm(ones_m) - (1 / n) * ones_E
        E = torch.clip(E, min=0, max=1)  # E,E.shape[0],E.shape[0],0)
        if i % 10 == 1:
            if torch.linalg.norm(E - prev_E) < ((10e-7) * n * n):
                break
    return E
//  [python]
def _compute_L1_penalty(self):
        """
        Compute L1 penalty.

        Return
        ------
        tf.Tensor: L1 penalty term (scalar-valued).
        """
        return torch.norm(self.B, p=1)
//  [python]
def _compute_L1_penalty(self):
        """
        Compute L1 penalty.

        Return
        ------
        tf.Tensor: L1 penalty term (scalar-valued).
        """
        return torch.linalg.norm(self.B, ord=1)
//  [python]
def modify_conformer_torsion_angles(pos, edge_index, mask_rotate, torsion_updates):
    for idx_edge, e in enumerate(edge_index):
        # if torsion_updates[idx_edge] == 0:
        #     continue
        u, v = e[0], e[1]

        # check if need to reverse the edge, v should be connected to the part that gets rotated
        # assert not mask_rotate[idx_edge, u]
        # assert mask_rotate[idx_edge, v]

        rot_vec = pos[u] - pos[v]  # convention: positive rotation if pointing inwards
        rot_vec = rot_vec / torch.norm(rot_vec)  # idx_edge!
        rot_mat = gen_matrix_from_rot_vec(rot_vec, torsion_updates[idx_edge])

        pos[mask_rotate[idx_edge]] = (pos[mask_rotate[idx_edge]] - pos[v]) @ rot_mat.T + pos[v]

    return pos
//  [python]
def modify_conformer_torsion_angles(pos, edge_index, mask_rotate, torsion_updates):
    for idx_edge, e in enumerate(edge_index):
        # if torsion_updates[idx_edge] == 0:
        #     continue
        u, v = e[0], e[1]

        # check if need to reverse the edge, v should be connected to the part that gets rotated
        # assert not mask_rotate[idx_edge, u]
        # assert mask_rotate[idx_edge, v]

        rot_vec = pos[u] - pos[v]  # convention: positive rotation if pointing inwards
        rot_vec = rot_vec / torch.linalg.norm(rot_vec)  # idx_edge!
        rot_mat = gen_matrix_from_rot_vec(rot_vec, torsion_updates[idx_edge])

        pos[mask_rotate[idx_edge]] = (pos[mask_rotate[idx_edge]] - pos[v]) @ rot_mat.T + pos[v]

    return pos
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    From https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html
    Convert rotations given as axis/angle to quaternions.
    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    From https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html
    Convert rotations given as axis/angle to quaternions.
    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
            torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
            0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def angle(v1: torch.Tensor, v2: torch.Tensor):
    """Compute angle between 2 vectors

    For robustness, we use the same formulation as in PPFNet, i.e.
        angle(v1, v2) = atan2(cross(v1, v2), dot(v1, v2)).
    This handles the case where one of the vectors is 0.0, since torch.atan2(0.0, 0.0)=0.0

    Args:
        v1: (B, *, 3)
        v2: (B, *, 3)

    Returns:

    """

    cross_prod = torch.stack([v1[..., 1] * v2[..., 2] - v1[..., 2] * v2[..., 1],
                              v1[..., 2] * v2[..., 0] - v1[..., 0] * v2[..., 2],
                              v1[..., 0] * v2[..., 1] - v1[..., 1] * v2[..., 0]], dim=-1)
    cross_prod_norm = torch.norm(cross_prod, dim=-1)
    dot_prod = torch.sum(v1 * v2, dim=-1)

    return torch.atan2(cross_prod_norm, dot_prod)
//  [python]
def angle(v1: torch.Tensor, v2: torch.Tensor):
    """Compute angle between 2 vectors

    For robustness, we use the same formulation as in PPFNet, i.e.
        angle(v1, v2) = atan2(cross(v1, v2), dot(v1, v2)).
    This handles the case where one of the vectors is 0.0, since torch.atan2(0.0, 0.0)=0.0

    Args:
        v1: (B, *, 3)
        v2: (B, *, 3)

    Returns:

    """

    cross_prod = torch.stack([v1[..., 1] * v2[..., 2] - v1[..., 2] * v2[..., 1],
                              v1[..., 2] * v2[..., 0] - v1[..., 0] * v2[..., 2],
                              v1[..., 0] * v2[..., 1] - v1[..., 1] * v2[..., 0]], dim=-1)
    cross_prod_norm = torch.linalg.norm(cross_prod, dim=-1)
    dot_prod = torch.sum(v1 * v2, dim=-1)

    return torch.atan2(cross_prod_norm, dot_prod)
//  [python]
def forward(self, data, clusts, edge_index, closest_index=None):

        # Check if the graph is undirected, select the relevant part of the edge index
        half_idx = int(edge_index.shape[1] / 2)
        undirected = not edge_index.shape[1] or (not edge_index.shape[1] % 2 and [edge_index[1, 0], edge_index[0, 0]] == edge_index[:, half_idx].tolist())
        if undirected: edge_index = edge_index[:, :half_idx]

        # If numpy is to be used, bring data to cpu, pass through Numba function
        # Otherwise use torch-based implementation of cluster_edge_features
        if self.use_numpy:
            feats = cluster_edge_features(data, clusts, edge_index.T, closest_index=closest_index)
        else:
            # Get the voxel set
            voxels = data[:, COORD_COLS].float()

            # Here is a torch-based implementation of cluster_edge_features
            feats = []
            for e in edge_index.T:

                # Get the voxels in the clusters connected by the edge
                x1 = voxels[clusts[e[0]]]
                x2 = voxels[clusts[e[1]]]

                # Find the closest set point in each cluster
                d12 = local_cdist(x1, x2)
                imin = torch.argmin(d12)
                i1, i2 = imin//len(x2), imin%len(x2)
                v1 = x1[i1,:] # closest point in c1
                v2 = x2[i2,:] # closest point in c2

                # Displacement
                disp = v1 - v2

                # Distance
                lend = torch.norm(disp)
                if lend > 0:
                    disp = disp / lend

                # Outer product
                B = torch.ger(disp, disp).flatten()

                feats.append(torch.cat([v1, v2, disp, lend.reshape(1), B]))

            feats = torch.stack(feats, dim=0)

        # If the graph is undirected, infer reciprocal features
        if undirected:
            feats_flip = feats.clone()
            feats_flip[:,:3] = feats[:,3:6]
            feats_flip[:,3:6] = feats[:,:3]
            feats_flip[:,6:9] = -feats[:,6:9]
            feats = torch.cat([feats,feats_flip])

        return feats
//  [python]
def forward(self, data, clusts, edge_index, closest_index=None):

        # Check if the graph is undirected, select the relevant part of the edge index
        half_idx = int(edge_index.shape[1] / 2)
        undirected = not edge_index.shape[1] or (not edge_index.shape[1] % 2 and [edge_index[1, 0], edge_index[0, 0]] == edge_index[:, half_idx].tolist())
        if undirected: edge_index = edge_index[:, :half_idx]

        # If numpy is to be used, bring data to cpu, pass through Numba function
        # Otherwise use torch-based implementation of cluster_edge_features
        if self.use_numpy:
            feats = cluster_edge_features(data, clusts, edge_index.T, closest_index=closest_index)
        else:
            # Get the voxel set
            voxels = data[:, COORD_COLS].float()

            # Here is a torch-based implementation of cluster_edge_features
            feats = []
            for e in edge_index.T:

                # Get the voxels in the clusters connected by the edge
                x1 = voxels[clusts[e[0]]]
                x2 = voxels[clusts[e[1]]]

                # Find the closest set point in each cluster
                d12 = local_cdist(x1, x2)
                imin = torch.argmin(d12)
                i1, i2 = imin//len(x2), imin%len(x2)
                v1 = x1[i1,:] # closest point in c1
                v2 = x2[i2,:] # closest point in c2

                # Displacement
                disp = v1 - v2

                # Distance
                lend = torch.linalg.norm(disp)
                if lend > 0:
                    disp = disp / lend

                # Outer product
                B = torch.ger(disp, disp).flatten()

                feats.append(torch.cat([v1, v2, disp, lend.reshape(1), B]))

            feats = torch.stack(feats, dim=0)

        # If the graph is undirected, infer reciprocal features
        if undirected:
            feats_flip = feats.clone()
            feats_flip[:,:3] = feats[:,3:6]
            feats_flip[:,3:6] = feats[:,:3]
            feats_flip[:,6:9] = -feats[:,6:9]
            feats = torch.cat([feats,feats_flip])

        return feats
//  [python]
def gradient_descent_full_line_torch(A, F, V, regu, gamma=1e-04, max_evals=250, verbosity=float('inf'), error=1e-07,
                                     regu_norm='L1', device='cpu'):
    """
    The function gradient_descent_full_line filters the signal by applying gradient descent method with line search.
    Parameters:
    A (numpy.ndarray): The gene expression matrix.
    F (numpy.ndarray): The filtering matrix.
    V (numpy.ndarray): The theoretic spectrum of covariance matrix.
    regu (float): The regularization coefficient.
    gamma (float, optional): The parameter for line search step size. Default is 1e-04.
    max_evals (int, optional): The maximum number of evaluations for the line search. Default is 250.
    verbosity (float, optional): The verbosity of the output. Default is float('inf').
    error (float, optional): The error tolerance for the line search. Default is 1e-07.
    regu_norm (str, optional): The type of regularization to apply. Default is 'L1'.

    Returns:
    numpy.ndarray: The filtering matrix.
    """
    VVT = V.mm(V.T)
    w = F
    evals = 0
    loss, grad = function_and_gradient_full_acc_torch(A=A, B=F, V=V, VVT=VVT, regu=regu, regu_norm=regu_norm)
    alpha = 1 / torch.norm(grad)
    prev_w = torch.zeros(w.shape, device=device)
    while evals < max_evals and torch.norm(w - prev_w) > error:
        prev_w = torch.clone(w)
        evals += 1
        if evals % verbosity == 0:
            print((evals))
            print('th Iteration    Loss :: ')
            print((loss))
        gTg = torch.norm(grad)
        gTg = gTg * gTg
        new_w = w - alpha * grad
        new_w = torch.clip(new_w, 0, 1)
        new_loss, new_grad = function_and_gradient_full_acc_torch(A=A, B=new_w, V=V,
                                                                  VVT=VVT, regu=regu, regu_norm=regu_norm)
        while new_loss > loss - gamma * alpha * gTg:
            alpha = ((alpha ** 2) * gTg) / (2 * (new_loss + alpha * gTg - loss))
            new_w = w - alpha * grad
            new_w = torch.clip(new_w, 0, 1)
            new_loss, new_grad = function_and_gradient_full_acc_torch(A=A, B=new_w, V=V, VVT=VVT,
                                                                      regu=regu, regu_norm=regu_norm)
        alpha = min(1, 2 * (loss - new_loss) / gTg)
        loss = new_loss
        grad = new_grad
        w = new_w
    del grad, loss, prev_w
    return w
//  [python]
def gradient_descent_full_line_torch(A, F, V, regu, gamma=1e-04, max_evals=250, verbosity=float('inf'), error=1e-07,
                                     regu_norm='L1', device='cpu'):
    """
    The function gradient_descent_full_line filters the signal by applying gradient descent method with line search.
    Parameters:
    A (numpy.ndarray): The gene expression matrix.
    F (numpy.ndarray): The filtering matrix.
    V (numpy.ndarray): The theoretic spectrum of covariance matrix.
    regu (float): The regularization coefficient.
    gamma (float, optional): The parameter for line search step size. Default is 1e-04.
    max_evals (int, optional): The maximum number of evaluations for the line search. Default is 250.
    verbosity (float, optional): The verbosity of the output. Default is float('inf').
    error (float, optional): The error tolerance for the line search. Default is 1e-07.
    regu_norm (str, optional): The type of regularization to apply. Default is 'L1'.

    Returns:
    numpy.ndarray: The filtering matrix.
    """
    VVT = V.mm(V.T)
    w = F
    evals = 0
    loss, grad = function_and_gradient_full_acc_torch(A=A, B=F, V=V, VVT=VVT, regu=regu, regu_norm=regu_norm)
    alpha = 1 / torch.linalg.norm(grad)
    prev_w = torch.zeros(w.shape, device=device)
    while evals < max_evals and torch.linalg.norm(w - prev_w) > error:
        prev_w = torch.clone(w)
        evals += 1
        if evals % verbosity == 0:
            print((evals))
            print('th Iteration    Loss :: ')
            print((loss))
        gTg = torch.linalg.norm(grad)
        gTg = gTg * gTg
        new_w = w - alpha * grad
        new_w = torch.clip(new_w, 0, 1)
        new_loss, new_grad = function_and_gradient_full_acc_torch(A=A, B=new_w, V=V,
                                                                  VVT=VVT, regu=regu, regu_norm=regu_norm)
        while new_loss > loss - gamma * alpha * gTg:
            alpha = ((alpha ** 2) * gTg) / (2 * (new_loss + alpha * gTg - loss))
            new_w = w - alpha * grad
            new_w = torch.clip(new_w, 0, 1)
            new_loss, new_grad = function_and_gradient_full_acc_torch(A=A, B=new_w, V=V, VVT=VVT,
                                                                      regu=regu, regu_norm=regu_norm)
        alpha = min(1, 2 * (loss - new_loss) / gTg)
        loss = new_loss
        grad = new_grad
        w = new_w
    del grad, loss, prev_w
    return w
//  [python]
def get_perpendicular_unit_vector(u):
    negative = torch.Tensor([1, -1]).to(u.device)
    # add dimensions to match u
    negative = torch.reshape(negative, [1] * (u.dim() - 1) + [2])
    u = u[..., [1, 0]] * negative
    u = u / torch.norm(u, dim=-1, keepdim=True)
    return u
//  [python]
def get_perpendicular_unit_vector(u):
    negative = torch.Tensor([1, -1]).to(u.device)
    # add dimensions to match u
    negative = torch.reshape(negative, [1] * (u.dim() - 1) + [2])
    u = u[..., [1, 0]] * negative
    u = u / torch.linalg.norm(u, dim=-1, keepdim=True)
    return u
//  [python]
def test_sample_from_unit_sphere_surface(num_dim: int, method: str):
    s = sample_from_hyper_sphere_surface(num_dim, method)
    assert 0.95 <= to.norm(s, p=2) <= 1.05
//  [python]
def test_sample_from_unit_sphere_surface(num_dim: int, method: str):
    s = sample_from_hyper_sphere_surface(num_dim, method)
    assert 0.95 <= torch.linalg.norm(s, p=2) <= 1.05
//  [python]
def convert_3d_to_uv_coordinates(points):
    # eps = 1E-8
    eps = 1E-4
    if type(points) == torch.Tensor:
        rad  = torch.clamp(torch.norm(points, p=2, dim=-1), min=eps)
        
        phi = torch.atan2(points[...,1], points[...,0])
        theta = torch.asin(torch.clamp(points[...,2]/rad, min=-1+eps, max=1-eps))
        u = 0.5 + phi/(2*np.pi)
        v = 0.5 + theta/np.pi
        return torch.stack([u,v],dim=-1)
    else:
        rad  = np.linalg.norm(points, axis=1)
        phi = np.arctan2(points[:,1], points[:,0])
        theta = np.arcsin(points[:,2]/rad)
        u = 0.5 + phi/(2*np.pi)
        v = 0.5 + theta/np.pi
        return np.stack([u,v],axis=1)
//  [python]
def convert_3d_to_uv_coordinates(points):
    # eps = 1E-8
    eps = 1E-4
    if type(points) == torch.Tensor:
        rad  = torch.clamp(torch.linalg.norm(points, p=2, dim=-1), min=eps)
        
        phi = torch.atan2(points[...,1], points[...,0])
        theta = torch.asin(torch.clamp(points[...,2]/rad, min=-1+eps, max=1-eps))
        u = 0.5 + phi/(2*np.pi)
        v = 0.5 + theta/np.pi
        return torch.stack([u,v],dim=-1)
    else:
        rad  = np.linalg.norm(points, axis=1)
        phi = np.arctan2(points[:,1], points[:,0])
        theta = np.arcsin(points[:,2]/rad)
        u = 0.5 + phi/(2*np.pi)
        v = 0.5 + theta/np.pi
        return np.stack([u,v],axis=1)
//  [python]
def raw_morlets(grid_or_shape, wave_vectors, gaussian_bases, morlet=True, ifftshift=True, fft=True):
    """ Helper function for creating morlet filters

        Parameters:
            grid_or_shape -- a grid of the size of the filter or a tuple that indicates its shape
            wave_vectors -- directions of the wave part of the morlet wavelet
            gaussian_bases -- bases of the gaussian part of the morlet wavelet
            morlet -- boolean for morlet or gabor wavelet
            ifftshift -- boolean for the ifftshift (inverse fast fourier transform shift)
            fft -- boolean for the fft (fast fourier transform)

        Returns:
            filters -- the wavelet filters before normalization
            
    """
    n_filters, n_dim = wave_vectors.shape
    assert gaussian_bases.shape == (n_filters, n_dim, n_dim)

    if isinstance(grid_or_shape, tuple):
        shape = grid_or_shape
        ranges = [torch.arange(-(s // 2), -(s // 2) + s, dtype=torch.float) for s in shape]
        grid = torch.stack(torch.meshgrid(*ranges), 0)
    else:
        shape = grid_or_shape.shape[1:]
        grid = grid_or_shape

    waves = torch.exp(1.0j * torch.matmul(grid.T, wave_vectors.T).T)
    gaussian_directions = torch.matmul(grid.T, gaussian_bases.T.reshape(n_dim, n_dim * n_filters)).T
    gaussian_directions = gaussian_directions.reshape((n_dim, n_filters) + shape)
    radii = torch.norm(gaussian_directions, dim=0)
    gaussians = torch.exp(-0.5 * radii ** 2)
    signal_dims = list(range(1, n_dim + 1))
    gabors = gaussians * waves

    if morlet:
        gaussian_sums = gaussians.sum(dim=signal_dims, keepdim=True)
        gabor_sums = gabors.sum(dim=signal_dims, keepdim=True).real
        morlets = gabors - gabor_sums / gaussian_sums * gaussians
        filters = morlets
    else:
        filters = gabors
    if ifftshift:
        filters = torch.fft.ifftshift(filters, dim=signal_dims)
    if fft:
        filters = torch.fft.fftn(filters, dim=signal_dims)

    return filters
//  [python]
def raw_morlets(grid_or_shape, wave_vectors, gaussian_bases, morlet=True, ifftshift=True, fft=True):
    """ Helper function for creating morlet filters

        Parameters:
            grid_or_shape -- a grid of the size of the filter or a tuple that indicates its shape
            wave_vectors -- directions of the wave part of the morlet wavelet
            gaussian_bases -- bases of the gaussian part of the morlet wavelet
            morlet -- boolean for morlet or gabor wavelet
            ifftshift -- boolean for the ifftshift (inverse fast fourier transform shift)
            fft -- boolean for the fft (fast fourier transform)

        Returns:
            filters -- the wavelet filters before normalization
            
    """
    n_filters, n_dim = wave_vectors.shape
    assert gaussian_bases.shape == (n_filters, n_dim, n_dim)

    if isinstance(grid_or_shape, tuple):
        shape = grid_or_shape
        ranges = [torch.arange(-(s // 2), -(s // 2) + s, dtype=torch.float) for s in shape]
        grid = torch.stack(torch.meshgrid(*ranges), 0)
    else:
        shape = grid_or_shape.shape[1:]
        grid = grid_or_shape

    waves = torch.exp(1.0j * torch.matmul(grid.T, wave_vectors.T).T)
    gaussian_directions = torch.matmul(grid.T, gaussian_bases.T.reshape(n_dim, n_dim * n_filters)).T
    gaussian_directions = gaussian_directions.reshape((n_dim, n_filters) + shape)
    radii = torch.linalg.norm(gaussian_directions, dim=0)
    gaussians = torch.exp(-0.5 * radii ** 2)
    signal_dims = list(range(1, n_dim + 1))
    gabors = gaussians * waves

    if morlet:
        gaussian_sums = gaussians.sum(dim=signal_dims, keepdim=True)
        gabor_sums = gabors.sum(dim=signal_dims, keepdim=True).real
        morlets = gabors - gabor_sums / gaussian_sums * gaussians
        filters = morlets
    else:
        filters = gabors
    if ifftshift:
        filters = torch.fft.ifftshift(filters, dim=signal_dims)
    if fft:
        filters = torch.fft.fftn(filters, dim=signal_dims)

    return filters
//  [python]
def gen_noisy_linear_weights(parameter_noise, weight):
    """Get Gaussian noisy linear weights based on given noise level ....
    and the weights themselves. The noise are normalized per channel (dim=1).

    InputArgs:
        parameter_noise: float, noise level, [0.0, 1.0]
        weight: Tensor, a weight tensor of a matrix
    Return:
        noisy_weight: Tensor, same dimension as weight, but with noise added.
    """
    noise = torch.randn_like(weight).to(device=weight.device)
    normalized_noise = noise / torch.norm(noise, dim=1, keepdim=True)

    w_norm = torch.norm(weight, dim=1, keepdim=True).detach()
    scale = parameter_noise * w_norm
    noisy_weight = weight + scale * normalized_noise
    return noisy_weight
//  [python]
def gen_noisy_linear_weights(parameter_noise, weight):
    """Get Gaussian noisy linear weights based on given noise level ....
    and the weights themselves. The noise are normalized per channel (dim=1).

    InputArgs:
        parameter_noise: float, noise level, [0.0, 1.0]
        weight: Tensor, a weight tensor of a matrix
    Return:
        noisy_weight: Tensor, same dimension as weight, but with noise added.
    """
    noise = torch.randn_like(weight).to(device=weight.device)
    normalized_noise = noise / torch.linalg.norm(noise, dim=1, keepdim=True)

    w_norm = torch.linalg.norm(weight, dim=1, keepdim=True).detach()
    scale = parameter_noise * w_norm
    noisy_weight = weight + scale * normalized_noise
    return noisy_weight
//  [python]
def get_face_normal(V, F):
    v1 = V[..., F[:, 1], :] - V[..., F[:, 0], :]
    v2 = V[..., F[:, 2], :] - V[..., F[:, 0], :]
    N = torch.linalg.cross(v1, v2)
    N = N / (torch.norm(N, dim=-1, keepdim=True) + 1e-8)
    return N
//  [python]
def get_face_normal(V, F):
    v1 = V[..., F[:, 1], :] - V[..., F[:, 0], :]
    v2 = V[..., F[:, 2], :] - V[..., F[:, 0], :]
    N = torch.linalg.cross(v1, v2)
    N = N / (torch.linalg.norm(N, dim=-1, keepdim=True) + 1e-8)
    return N
//  [python]
def weighted_8points(x_in, logits):
    # 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)
    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):
    # 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)
    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 get_local_frames(V, F):
    if F.dtype != torch.int64:
        F = F.to(torch.int64)
    v1 = V[..., F[:, 1], :] - V[..., F[:, 0], :]
    v2 = V[..., F[:, 2], :] - V[..., F[:, 0], :]
    N = torch.linalg.cross(v1, v2)
    assert torch.allclose(torch.linalg.cross(v1, v2)[0], torch.cross(v1[0], v2[0]))
    N = N / (torch.norm(N, dim=-1, keepdim=True) + 1e-8)
    local_frames = torch.stack([v1, v2, N], dim=-1)
    return local_frames
//  [python]
def get_local_frames(V, F):
    if F.dtype != torch.int64:
        F = F.to(torch.int64)
    v1 = V[..., F[:, 1], :] - V[..., F[:, 0], :]
    v2 = V[..., F[:, 2], :] - V[..., F[:, 0], :]
    N = torch.linalg.cross(v1, v2)
    assert torch.allclose(torch.linalg.cross(v1, v2)[0], torch.cross(v1[0], v2[0]))
    N = N / (torch.linalg.norm(N, dim=-1, keepdim=True) + 1e-8)
    local_frames = torch.stack([v1, v2, N], dim=-1)
    return local_frames