jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def step_correct(
        self,
        model_output: torch.FloatTensor,
        sample: torch.FloatTensor,
        generator: Optional[torch.Generator] = None,
        return_dict: bool = True,
    ) -> Union[SchedulerOutput, Tuple]:
        """
        Correct the predicted sample based on the `model_output` of the network. This is often run repeatedly after
        making the prediction for the previous timestep.

        Args:
            model_output (`torch.FloatTensor`):
                The direct output from learned diffusion model.
            sample (`torch.FloatTensor`):
                A current instance of a sample created by the diffusion process.
            generator (`torch.Generator`, *optional*):
                A random number generator.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`.

        Returns:
            [`~schedulers.scheduling_sde_ve.SdeVeOutput`] or `tuple`:
                If return_dict is `True`, [`~schedulers.scheduling_sde_ve.SdeVeOutput`] is returned, otherwise a tuple
                is returned where the first element is the sample tensor.

        """
        if self.timesteps is None:
            raise ValueError(
                "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"
            )

        # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
        # sample noise for correction
        noise = randn_tensor(sample.shape, layout=sample.layout, generator=generator).to(sample.device)

        # compute step size from the model_output, the noise, and the snr
        grad_norm = torch.linalg.norm(model_output.reshape(model_output.shape[0], -1), dim=-1).mean()
        noise_norm = torch.linalg.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
        step_size = (self.config.snr * noise_norm / grad_norm) ** 2 * 2
        step_size = step_size * torch.ones(sample.shape[0]).to(sample.device)
        # self.repeat_scalar(step_size, sample.shape[0])

        # compute corrected sample: model_output term and noise term
        step_size = step_size.flatten()
        while len(step_size.shape) < len(sample.shape):
            step_size = step_size.unsqueeze(-1)
        prev_sample_mean = sample + step_size * model_output
        prev_sample = prev_sample_mean + ((step_size * 2) ** 0.5) * noise

        if not return_dict:
            return (prev_sample,)

        return SchedulerOutput(prev_sample=prev_sample)
//  [python]
def update_converged_count(self):
        """Determine the number of converged eigenpairs using backward stable
        convergence criterion, see discussion in Sec 4.3 of [DuerschEtal2018].
        Users may redefine this method for custom convergence criteria.
        """
        # (...) -> int
        prev_count = self.ivars['converged_count']
        tol = self.fparams['tol']
        A_norm = self.fvars['A_norm']
        B_norm = self.fvars['B_norm']
        E, X, R = self.E, self.X, self.R
        rerr = torch.norm(R, 2, (0, )) * (torch.norm(X, 2, (0, )) * (A_norm + E[:X.shape[-1]] * B_norm)) ** -1
        converged = rerr < tol
        # print(rerr)
        # print(rerr.mean())
        count = 0
        for b in converged:
            if not b:
                # ignore convergence of following pairs to ensure
                # strict ordering of eigenpairs
                break
            count += 1
        assert count >= prev_count, 'the number of converged eigenpairs ' \
            '(was {}, got {}) cannot decrease'.format(prev_count, count)
        self.ivars['converged_count'] = count
        self.tvars['rerr'] = rerr
        return count
//  [python]
def update_converged_count(self):
        """Determine the number of converged eigenpairs using backward stable
        convergence criterion, see discussion in Sec 4.3 of [DuerschEtal2018].
        Users may redefine this method for custom convergence criteria.
        """
        # (...) -> int
        prev_count = self.ivars['converged_count']
        tol = self.fparams['tol']
        A_norm = self.fvars['A_norm']
        B_norm = self.fvars['B_norm']
        E, X, R = self.E, self.X, self.R
        rerr = torch.linalg.norm(R, 2, (0, )) * (torch.linalg.norm(X, 2, (0, )) * (A_norm + E[:X.shape[-1]] * B_norm)) ** -1
        converged = rerr < tol
        # print(rerr)
        # print(rerr.mean())
        count = 0
        for b in converged:
            if not b:
                # ignore convergence of following pairs to ensure
                # strict ordering of eigenpairs
                break
            count += 1
        assert count >= prev_count, 'the number of converged eigenpairs ' \
            '(was {}, got {}) cannot decrease'.format(prev_count, count)
        self.ivars['converged_count'] = count
        self.tvars['rerr'] = rerr
        return count
//  [python]
def get_spatial_compatibility(self, infos):
        infos["src_dist"] = torch.norm((infos["src_keypts"][:, :, None, :] - infos["src_keypts"][:, None, :, :]), dim=-1)
        infos["tgt_dist"] = torch.norm((infos["tgt_keypts"][:, :, None, :] - infos["tgt_keypts"][:, None, :, :]), dim=-1)
        spatial_compatibility = infos["src_dist"] - infos["tgt_dist"]
        infos["spatial_compatibility"] = torch.clamp(1.0 - spatial_compatibility ** 2 / self.sigma_spat ** 2, min=0)
        return infos
//  [python]
def get_spatial_compatibility(self, infos):
        infos["src_dist"] = torch.linalg.norm((infos["src_keypts"][:, :, None, :] - infos["src_keypts"][:, None, :, :]), dim=-1)
        infos["tgt_dist"] = torch.linalg.norm((infos["tgt_keypts"][:, :, None, :] - infos["tgt_keypts"][:, None, :, :]), dim=-1)
        spatial_compatibility = infos["src_dist"] - infos["tgt_dist"]
        infos["spatial_compatibility"] = torch.clamp(1.0 - spatial_compatibility ** 2 / self.sigma_spat ** 2, min=0)
        return infos
//  [python]
def eval_trans(self, infos):
        seedwise_trans = infos["seedwise_trans"]  # [B, K, 4, 4]
        pred_position = torch.einsum('bsnm,bmk->bsnk', seedwise_trans[:, :, :3, :3], infos["src_keypts"].permute(0, 2, 1)) + seedwise_trans[:, :, :3, 3:4]  # [bs, num_seeds, num_corr, 3]
        pred_position = pred_position.permute(0, 1, 3, 2)
        infos["L2_dis"] = torch.norm(pred_position - infos["tgt_keypts"][:, None, :, :], dim=-1)  # [B, K, N]
        infos = self.metric(infos)
        infos["batch_best_guess"] = infos["seedwise_fitness"].argmax(dim=1)  # [B]
        return infos
//  [python]
def eval_trans(self, infos):
        seedwise_trans = infos["seedwise_trans"]  # [B, K, 4, 4]
        pred_position = torch.einsum('bsnm,bmk->bsnk', seedwise_trans[:, :, :3, :3], infos["src_keypts"].permute(0, 2, 1)) + seedwise_trans[:, :, :3, 3:4]  # [bs, num_seeds, num_corr, 3]
        pred_position = pred_position.permute(0, 1, 3, 2)
        infos["L2_dis"] = torch.linalg.norm(pred_position - infos["tgt_keypts"][:, None, :, :], dim=-1)  # [B, K, N]
        infos = self.metric(infos)
        infos["batch_best_guess"] = infos["seedwise_fitness"].argmax(dim=1)  # [B]
        return infos
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, it_num, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4]
            - src_keypts:    [bs, num_corr, 3]
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [K] 
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        inlier_threshold = 1.2

        # inlier_threshold_list = [self.inlier_threshold] * it_num

        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, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
                # weights=((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, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4]
            - src_keypts:    [bs, num_corr, 3]
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [K] 
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        inlier_threshold = 1.2

        # inlier_threshold_list = [self.inlier_threshold] * it_num

        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, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def _construct_model(
        device: torch.device,
    ) -> Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]:
        data = torch.arange(80, dtype=torch.float, device=device)
        data /= torch.norm(data)
        model = nn.Sequential(
            nn.Linear(80, 40, bias=False),
            nn.Linear(40, 1, bias=False),
        ).to(device=device)
        model[0].weight.data.fill_(0.01)
        model[1].weight.data.fill_(0.01)
        model[1].requires_grad = False
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=device)
        return model, loss, data, target
//  [python]
def _construct_model(
        device: torch.device,
    ) -> Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]:
        data = torch.arange(80, dtype=torch.float, device=device)
        data /= torch.linalg.norm(data)
        model = nn.Sequential(
            nn.Linear(80, 40, bias=False),
            nn.Linear(40, 1, bias=False),
        ).to(device=device)
        model[0].weight.data.fill_(0.01)
        model[1].weight.data.fill_(0.01)
        model[1].requires_grad = False
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=device)
        return model, loss, data, target
//  [python]
def Lasso_sigmoid(matrix, t):
        """
        A trick that can make use of both the value of bool(lasso > t) and the gradient of sigmoid(lasso - t)

        Parameters
        ----------
        matrix : Tensor
            the matrix to calculate lasso norm
        t : float
            the threshold
        """
        lasso = torch.norm(matrix) - t
        indicator = (lasso > 0).float()  # torch.sign(lasso)
        with torch.no_grad():
            #            indicator = indicator / 2 + .5 # realign indicator from (-1, 1) to (0, 1)
            indicator -= F.sigmoid(lasso)
        indicator += F.sigmoid(lasso)
        return indicator
//  [python]
def Lasso_sigmoid(matrix, t):
        """
        A trick that can make use of both the value of bool(lasso > t) and the gradient of sigmoid(lasso - t)

        Parameters
        ----------
        matrix : Tensor
            the matrix to calculate lasso norm
        t : float
            the threshold
        """
        lasso = torch.linalg.norm(matrix) - t
        indicator = (lasso > 0).float()  # torch.sign(lasso)
        with torch.no_grad():
            #            indicator = indicator / 2 + .5 # realign indicator from (-1, 1) to (0, 1)
            indicator -= F.sigmoid(lasso)
        indicator += F.sigmoid(lasso)
        return indicator
//  [python]
def weighted_8points(coords, logits):
  # logits shape = (batch, num_point, 1)
  # coords shape = (batch, 4, num_point)
  w = torch.nn.functional.relu(torch.tanh(logits))
  X = torch.stack([
      coords[:, 2] * coords[:, 0], coords[:, 2] * coords[:, 1], coords[:, 2],
      coords[:, 3] * coords[:, 0], coords[:, 3] * coords[:, 1], coords[:, 3],
      coords[:, 0], coords[:, 1],
      torch.ones_like(coords[:, 0])
  ],
                  dim=1).transpose(2, 1)
  # wX shape = (batch, num_point, 9)
  wX = w.unsqueeze(-1) * X
  # XwX shape = (batch, 9, 9)
  XwX = torch.bmm(X.transpose(2, 1), wX)

  v = batch_symeig(XwX)
  # _, v = torch.symeig(XwX, eigenvectors=True)
  e = torch.reshape(v[:, :, 0], (logits.shape[0], 9))
  e = e / torch.norm(e, dim=1, keepdim=True)
  return e
//  [python]
def weighted_8points(coords, logits):
  # logits shape = (batch, num_point, 1)
  # coords shape = (batch, 4, num_point)
  w = torch.nn.functional.relu(torch.tanh(logits))
  X = torch.stack([
      coords[:, 2] * coords[:, 0], coords[:, 2] * coords[:, 1], coords[:, 2],
      coords[:, 3] * coords[:, 0], coords[:, 3] * coords[:, 1], coords[:, 3],
      coords[:, 0], coords[:, 1],
      torch.ones_like(coords[:, 0])
  ],
                  dim=1).transpose(2, 1)
  # wX shape = (batch, num_point, 9)
  wX = w.unsqueeze(-1) * X
  # XwX shape = (batch, 9, 9)
  XwX = torch.bmm(X.transpose(2, 1), wX)

  v = batch_symeig(XwX)
  # _, v = torch.symeig(XwX, eigenvectors=True)
  e = torch.reshape(v[:, :, 0], (logits.shape[0], 9))
  e = e / torch.linalg.norm(e, dim=1, keepdim=True)
  return e
//  [python]
def l2_norm(input, axis=1):
	norm = torch.norm(input, 2, axis, True)
	output = torch.div(input, norm)
	return output
//  [python]
def l2_norm(input, axis=1):
	norm = torch.linalg.norm(input, 2, axis, keepdims=True)
	output = torch.div(input, norm)
	return output
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        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] * 20
        else: # for KITTI
            inlier_threshold_list = [1.2] * 20

        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, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1/(1 + (L2_dis/inlier_threshold)**2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        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] * 20
        else: # for KITTI
            inlier_threshold_list = [1.2] * 20

        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, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1/(1 + (L2_dis/inlier_threshold)**2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def batch_rodrigues(axisang):
    #axisang N x 3
    axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)
    rot_mat = quat2mat(quat)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def batch_rodrigues(axisang):
    #axisang N x 3
    axisang_norm = torch.linalg.norm(axisang + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)
    rot_mat = quat2mat(quat)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def SC2_PCR(self, src_keypts, tgt_keypts):
        """
        Input:
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
        Output:
            - potential_trans_by_IC:   [bs, 4, 4], the best transformation matrix selected by IC metric.
            - best_trans_by_IC:  [bs, N_s^{'} 4, 4], the potential transformation matrix selected by IC metric.
        """
        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
        #################################
        potential_trans_by_IC, best_trans_by_IC = self.cal_seed_trans(seeds, SC2_measure, src_keypts, tgt_keypts)

        return potential_trans_by_IC, best_trans_by_IC
//  [python]
def SC2_PCR(self, src_keypts, tgt_keypts):
        """
        Input:
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
        Output:
            - potential_trans_by_IC:   [bs, 4, 4], the best transformation matrix selected by IC metric.
            - best_trans_by_IC:  [bs, N_s^{'} 4, 4], the potential transformation matrix selected by IC metric.
        """
        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
        #################################
        potential_trans_by_IC, best_trans_by_IC = self.cal_seed_trans(seeds, SC2_measure, src_keypts, tgt_keypts)

        return potential_trans_by_IC, best_trans_by_IC
//  [python]
def forward(self, x):
        x = self.nn(x.max(dim=1)[0])
        rotation = self.proj_rot(x)
        translation = self.proj_trans(x)
        rotation = rotation / torch.norm(rotation, p=2, dim=1, keepdim=True)
        return rotation, translation
//  [python]
def forward(self, x):
        x = self.nn(x.max(dim=1)[0])
        rotation = self.proj_rot(x)
        translation = self.proj_trans(x)
        rotation = rotation / torch.linalg.norm(rotation, ord=2, dim=1, keepdim=True)
        return rotation, translation
//  [python]
def _safe_norm(v):
    if not torch.isfinite(v).all():
        return torch.tensor(float("inf"), dtype=v.dtype, device=v.device)
    return torch.norm(v)
//  [python]
def _safe_norm(v):
    if not torch.isfinite(v).all():
        return torch.tensor(float("inf"), dtype=v.dtype, device=v.device)
    return torch.linalg.norm(v)
//  [python]
def fr_distance_univariate_gaussian(
    mu_1: torch.Tensor, sig_1: torch.Tensor, mu_2: torch.Tensor, sig_2: torch.Tensor
) -> torch.Tensor:
    """Calculates the Fisher-Rao distance between univariate gaussian distributions in prallel.

    Args:
        mu_1 (torch.Tensor): Tensor of dimension (N,*) containing the means of N different univariate gaussians
        sig_1 (torch.Tensor): Standard deviations of univariate gaussian distributions
        mu_2 (torch.Tensor): Means of the second univariate gaussian distributions
        sig_2 (torch.Tensor): Standard deviation of the second univariate gaussian distributions
    Returns:
        torch.Tensor: Distance tensor of size (N,*)
    """
    dim = len(mu_1.shape)
    mu_1, mu_2 = mu_1.reshape(*mu_1.shape, 1), mu_2.reshape(*mu_2.shape, 1)
    sig_1, sig_2 = sig_1.reshape(*sig_1.shape, 1), sig_2.reshape(*sig_2.shape, 1)

    sqrt_2 = torch.sqrt(torch.tensor(2.0, device=mu_1.device))
    a = torch.norm(
        torch.cat((mu_1 / sqrt_2, sig_1), dim=dim)
        - torch.cat((mu_2 / sqrt_2, -1 * sig_2), dim=dim),
        p=2,
        dim=dim,
    )
    b = torch.norm(
        torch.cat((mu_1 / sqrt_2, sig_1), dim=dim)
        - torch.cat((mu_2 / sqrt_2, sig_2), dim=dim),
        p=2,
        dim=dim,
    )

    num = a + b + 1e-12
    den = a - b + 1e-12
    return sqrt_2 * torch.log(num / den)
//  [python]
def fr_distance_univariate_gaussian(
    mu_1: torch.Tensor, sig_1: torch.Tensor, mu_2: torch.Tensor, sig_2: torch.Tensor
) -> torch.Tensor:
    """Calculates the Fisher-Rao distance between univariate gaussian distributions in prallel.

    Args:
        mu_1 (torch.Tensor): Tensor of dimension (N,*) containing the means of N different univariate gaussians
        sig_1 (torch.Tensor): Standard deviations of univariate gaussian distributions
        mu_2 (torch.Tensor): Means of the second univariate gaussian distributions
        sig_2 (torch.Tensor): Standard deviation of the second univariate gaussian distributions
    Returns:
        torch.Tensor: Distance tensor of size (N,*)
    """
    dim = len(mu_1.shape)
    mu_1, mu_2 = mu_1.reshape(*mu_1.shape, 1), mu_2.reshape(*mu_2.shape, 1)
    sig_1, sig_2 = sig_1.reshape(*sig_1.shape, 1), sig_2.reshape(*sig_2.shape, 1)

    sqrt_2 = torch.sqrt(torch.tensor(2.0, device=mu_1.device))
    a = torch.linalg.norm(
        torch.cat((mu_1 / sqrt_2, sig_1), dim=dim)
        - torch.cat((mu_2 / sqrt_2, -1 * sig_2), dim=dim),
        ord=2,
        dim=dim,
    )
    b = torch.linalg.norm(
        torch.cat((mu_1 / sqrt_2, sig_1), dim=dim)
        - torch.cat((mu_2 / sqrt_2, sig_2), dim=dim),
        ord=2,
        dim=dim,
    )

    num = a + b + 1e-12
    den = a - b + 1e-12
    return sqrt_2 * torch.log(num / den)
//  [python]
def select_experts(
    gate_logits: torch.Tensor, top_k: int, moe_normalize_expert_weights: int
):
    # all_probs: (sequence_length, n_experts) and upcast for softmax
    all_probs = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)
    # weights, selected_experts: (sequence_length, top-k)
    weights, selected_experts = torch.topk(all_probs, top_k, dim=-1)
    if moe_normalize_expert_weights:
        weights = weights / torch.norm(
            weights, p=moe_normalize_expert_weights, dim=-1, keepdim=True
        )
    weights = weights.view(-1)
    selected_experts = selected_experts.view(-1)

    return selected_experts, weights
//  [python]
def select_experts(
    gate_logits: torch.Tensor, top_k: int, moe_normalize_expert_weights: int
):
    # all_probs: (sequence_length, n_experts) and upcast for softmax
    all_probs = torch.nn.functional.softmax(gate_logits, dim=1, dtype=torch.float)
    # weights, selected_experts: (sequence_length, top-k)
    weights, selected_experts = torch.topk(all_probs, top_k, dim=-1)
    if moe_normalize_expert_weights:
        weights = weights / torch.linalg.norm(
            weights, ord=moe_normalize_expert_weights, dim=-1, keepdim=True
        )
    weights = weights.view(-1)
    selected_experts = selected_experts.view(-1)

    return selected_experts, weights
//  [python]
def forward(self, x):
        agg_norm = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
        stand_div_norm = agg_norm / (agg_norm.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * stand_div_norm) + self.beta + x
//  [python]
def forward(self, x):
        agg_norm = torch.linalg.norm(x, ord=2, dim=(1, 2), keepdim=True)
        stand_div_norm = agg_norm / (agg_norm.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * stand_div_norm) + self.beta + x
//  [python]
def phaseless_loss(target, predicted, dim=-1, eps=1e-6):
    # TODO: Move sum and logging out of loss
    norm_target = torch.norm(target, 2, dim, keepdim=True)
    norm_predicted = torch.norm(predicted, 2, dim, keepdim=True)

    loss_length = (norm_target - norm_predicted) ** 2
    loss_length = torch.mean(loss_length.view(-1))

    overlap = torch.sum(target * predicted / (norm_target * norm_predicted + eps), dim)
    loss_overlap = 1.0 - overlap ** 2
    loss_overlap = torch.mean(loss_overlap.view(-1))

    #print(loss_length, loss_overlap, 'length/overlap')

    loss = loss_length + loss_overlap
    return loss
//  [python]
def phaseless_loss(target, predicted, dim=-1, eps=1e-6):
    # TODO: Move sum and logging out of loss
    norm_target = torch.linalg.norm(target, 2, dim, keepdim=True)
    norm_predicted = torch.linalg.norm(predicted, 2, dim, keepdim=True)

    loss_length = (norm_target - norm_predicted) ** 2
    loss_length = torch.mean(loss_length.view(-1))

    overlap = torch.sum(target * predicted / (norm_target * norm_predicted + eps), dim)
    loss_overlap = 1.0 - overlap ** 2
    loss_overlap = torch.mean(loss_overlap.view(-1))

    #print(loss_length, loss_overlap, 'length/overlap')

    loss = loss_length + loss_overlap
    return loss
//  [python]
def norm(tensor, order=None, axis=None):
        # pytorch does not accept `None` for any keyword arguments. additionally,
        # pytorch doesn't seems to support keyword arguments in the first place
        kwds = {}
        if axis is not None:
            kwds["dim"] = axis
        if order and order != "inf":
            kwds["p"] = order

        if order == "inf":
            res = torch.max(torch.abs(tensor), **kwds)
            if axis is not None:
                return res[0]  # ignore indices output
            return res
        return torch.norm(tensor, **kwds)
//  [python]
def norm(tensor, order=None, axis=None):
        # pytorch does not accept `None` for any keyword arguments. additionally,
        # pytorch doesn't seems to support keyword arguments in the first place
        kwds = {}
        if axis is not None:
            kwds["dim"] = axis
        if order and order != "inf":
            kwds["ord"] = order

        if order == "inf":
            res = torch.max(torch.abs(tensor), **kwds)
            if axis is not None:
                return res[0]  # ignore indices output
            return res
        return torch.linalg.norm(tensor, **kwds)
//  [python]
def exceeds_motion_thresholds(current_c2w: torch.Tensor, last_submap_c2w: torch.Tensor,
                              rot_thre: float = 50, trans_thre: float = 0.5) -> bool:
    """  Checks if a camera motion exceeds certain rotation and translation thresholds
    Args:
        current_c2w: The current camera-to-world transformation matrix.
        last_submap_c2w: The last submap's camera-to-world transformation matrix.
        rot_thre: The rotation threshold for triggering a new submap.
        trans_thre: The translation threshold for triggering a new submap.

    Returns:
        A boolean indicating whether a new submap is required.
    """
    delta_pose = torch.matmul(torch.linalg.inv(last_submap_c2w).float(), current_c2w.float())
    translation_diff = torch.norm(delta_pose[:3, 3])
    rot_euler_diff_deg = torch.abs(rotation_to_euler(delta_pose[:3, :3]))
    exceeds_thresholds = (translation_diff > trans_thre) or torch.any(rot_euler_diff_deg > rot_thre)
    return exceeds_thresholds.item()
//  [python]
def exceeds_motion_thresholds(current_c2w: torch.Tensor, last_submap_c2w: torch.Tensor,
                              rot_thre: float = 50, trans_thre: float = 0.5) -> bool:
    """  Checks if a camera motion exceeds certain rotation and translation thresholds
    Args:
        current_c2w: The current camera-to-world transformation matrix.
        last_submap_c2w: The last submap's camera-to-world transformation matrix.
        rot_thre: The rotation threshold for triggering a new submap.
        trans_thre: The translation threshold for triggering a new submap.

    Returns:
        A boolean indicating whether a new submap is required.
    """
    delta_pose = torch.matmul(torch.linalg.inv(last_submap_c2w).float(), current_c2w.float())
    translation_diff = torch.linalg.norm(delta_pose[:3, 3])
    rot_euler_diff_deg = torch.abs(rotation_to_euler(delta_pose[:3, :3]))
    exceeds_thresholds = (translation_diff > trans_thre) or torch.any(rot_euler_diff_deg > rot_thre)
    return exceeds_thresholds.item()
//  [python]
def min_loss_single_old(target, predicted, smooth=True, smooth_nonvec=False, loss_length=True):
    diff_a = (target - predicted) ** 2
    diff_b = (target + predicted) ** 2

    # Correct phase loss
    if smooth:
        reduce_dims = predicted.shape[:2] + (-1,)
        diff_a = diff_a.view(reduce_dims)
        diff_b = diff_b.view(reduce_dims)
        a = torch.mean(diff_a, 2, keepdim=True)
        b = torch.mean(diff_b, 2, keepdim=True)
        z = a + b + 1e-6
        coeff_a = b / z
        coeff_b = a / z
        diff = coeff_a * diff_a + coeff_b * diff_b
        #print(torch.mean(coeff_a), torch.mean(coeff_b), 'A/B')
    else:
        diff = torch.min(diff_a,diff_b)

    if smooth_nonvec:
        # don't take the mean of non-vectorial properties - those are from different states
        z = diff_a + diff_b + 1e-6
        coeff_a = diff_b / z
        coeff_b = diff_a / z
        diff = coeff_a * diff_a + coeff_b * diff_b
    loss = torch.mean(diff.view(-1))

    if loss_length:
        norm_target = torch.norm(target, 2, -1, keepdim=True)
        norm_predicted = torch.norm(predicted, 2, -1, keepdim=True)
        loss_length = (norm_target - norm_predicted) ** 2
        loss_length = torch.mean(loss_length.view(-1))
        #print(loss_length, 'LL')
        loss = loss + loss_length

    return loss
//  [python]
def min_loss_single_old(target, predicted, smooth=True, smooth_nonvec=False, loss_length=True):
    diff_a = (target - predicted) ** 2
    diff_b = (target + predicted) ** 2

    # Correct phase loss
    if smooth:
        reduce_dims = predicted.shape[:2] + (-1,)
        diff_a = diff_a.view(reduce_dims)
        diff_b = diff_b.view(reduce_dims)
        a = torch.mean(diff_a, 2, keepdim=True)
        b = torch.mean(diff_b, 2, keepdim=True)
        z = a + b + 1e-6
        coeff_a = b / z
        coeff_b = a / z
        diff = coeff_a * diff_a + coeff_b * diff_b
        #print(torch.mean(coeff_a), torch.mean(coeff_b), 'A/B')
    else:
        diff = torch.min(diff_a,diff_b)

    if smooth_nonvec:
        # don't take the mean of non-vectorial properties - those are from different states
        z = diff_a + diff_b + 1e-6
        coeff_a = diff_b / z
        coeff_b = diff_a / z
        diff = coeff_a * diff_a + coeff_b * diff_b
    loss = torch.mean(diff.view(-1))

    if loss_length:
        norm_target = torch.linalg.norm(target, 2, -1, keepdim=True)
        norm_predicted = torch.linalg.norm(predicted, 2, -1, keepdim=True)
        loss_length = (norm_target - norm_predicted) ** 2
        loss_length = torch.mean(loss_length.view(-1))
        #print(loss_length, 'LL')
        loss = loss + loss_length

    return loss
//  [python]
def forward(self, x):
        x_norm = torch.norm(x, p=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + 1e-05)
        if not self.class_wise_learnable_norm:
            fc_norm = torch.norm(self.fc.weight.data, p=2, dim=1).unsqueeze(1).expand_as(self.fc.weight.data)
            self.fc.weight.data = self.fc.weight.data.div(fc_norm + 1e-05)
        cos_dist = self.fc(x_normalized)
        score = self.scale_factor * cos_dist
        return score
//  [python]
def forward(self, x):
        x_norm = torch.linalg.norm(x, ord=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + 1e-05)
        if not self.class_wise_learnable_norm:
            fc_norm = torch.linalg.norm(self.fc.weight.data, ord=2, dim=1).unsqueeze(1).expand_as(self.fc.weight.data)
            self.fc.weight.data = self.fc.weight.data.div(fc_norm + 1e-05)
        cos_dist = self.fc(x_normalized)
        score = self.scale_factor * cos_dist
        return score
//  [python]
def intersect_sphere(ray_o, ray_d, radius=1):
    '''
    ray_o, ray_d: [..., 3]
    compute the depth of the intersection point between this ray and unit sphere
    '''
    # note: d1 becomes negative if this mid point is behind camera
    d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
    p = ray_o + d1.unsqueeze(-1) * ray_d
    # consider the case where the ray does not intersect the sphere
    ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
    d2 = radius ** 2 - torch.sum(p * p, dim=-1)
    mask = (d2 > 0)
    d2 = torch.sqrt(d2.clamp(min=1e-6)) * ray_d_cos
    d1, d2 = d1.unsqueeze(-1), d2.unsqueeze(-1)
    depth_range = [d1 - d2, d1 + d2]
    return depth_range, mask
//  [python]
def intersect_sphere(ray_o, ray_d, radius=1):
    '''
    ray_o, ray_d: [..., 3]
    compute the depth of the intersection point between this ray and unit sphere
    '''
    # note: d1 becomes negative if this mid point is behind camera
    d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
    p = ray_o + d1.unsqueeze(-1) * ray_d
    # consider the case where the ray does not intersect the sphere
    ray_d_cos = 1. / torch.linalg.norm(ray_d, dim=-1)
    d2 = radius ** 2 - torch.sum(p * p, dim=-1)
    mask = (d2 > 0)
    d2 = torch.sqrt(d2.clamp(min=1e-6)) * ray_d_cos
    d1, d2 = d1.unsqueeze(-1), d2.unsqueeze(-1)
    depth_range = [d1 - d2, d1 + d2]
    return depth_range, mask
//  [python]
def forward(self, user, item, neg_item=None):
        if neg_item is not None:  # used for train
            user_bais = self.user_bais(user).squeeze(-1)
            item_bais = self.item_bais(item).squeeze(-1)
            neg_item_bais = self.item_bais(neg_item).squeeze(-1)

            userembedding = self.user_laten(user)
            itemembedding = self.item_laten(item)
            neg_itemembedding = self.item_laten(neg_item)

            interaction = torch.mul(userembedding,itemembedding).sum(dim=-1)
            neg_interaction = torch.mul(userembedding,neg_itemembedding).sum(dim=-1)
            result_pos = user_bais + item_bais + interaction
            result_neg = user_bais + neg_item_bais + neg_interaction
            score = result_pos -result_neg
            bpr_loss = -torch.sum(F.logsigmoid(score))

            l2loss = torch.norm(userembedding, dim=-1).sum() + torch.norm(itemembedding,dim=-1).sum() + torch.norm(neg_itemembedding).sum()
            return bpr_loss,l2loss

        else:         # used for test
            user_bais = self.user_bais(user).squeeze(-1)
            item_bais = self.item_bais(item).squeeze(-1)
            userembedding = self.user_laten(user)
            itemembedding = self.item_laten(item)
            interaction = torch.mul(userembedding, itemembedding).sum(dim=-1)
            result = user_bais + item_bais + interaction
            return userembedding, itemembedding, result
//  [python]
def forward(self, user, item, neg_item=None):
        if neg_item is not None:  # used for train
            user_bais = self.user_bais(user).squeeze(-1)
            item_bais = self.item_bais(item).squeeze(-1)
            neg_item_bais = self.item_bais(neg_item).squeeze(-1)

            userembedding = self.user_laten(user)
            itemembedding = self.item_laten(item)
            neg_itemembedding = self.item_laten(neg_item)

            interaction = torch.mul(userembedding,itemembedding).sum(dim=-1)
            neg_interaction = torch.mul(userembedding,neg_itemembedding).sum(dim=-1)
            result_pos = user_bais + item_bais + interaction
            result_neg = user_bais + neg_item_bais + neg_interaction
            score = result_pos -result_neg
            bpr_loss = -torch.sum(F.logsigmoid(score))

            l2loss = torch.linalg.norm(userembedding, dim=-1).sum() + torch.linalg.norm(itemembedding,dim=-1).sum() + torch.linalg.norm(neg_itemembedding).sum()
            return bpr_loss,l2loss

        else:         # used for test
            user_bais = self.user_bais(user).squeeze(-1)
            item_bais = self.item_bais(item).squeeze(-1)
            userembedding = self.user_laten(user)
            itemembedding = self.item_laten(item)
            interaction = torch.mul(userembedding, itemembedding).sum(dim=-1)
            result = user_bais + item_bais + interaction
            return userembedding, itemembedding, result
//  [python]
def test_norm(self, device):
        # full reduction
        x = torch.randn(25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = x.norm(p).item()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res, expected, atol=1e-5, rtol=0, msg="full reduction failed for {}-norm".format(p))

        # one dimension
        x = torch.randn(25, 25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = x.norm(p, 1).cpu()
            expected = np.linalg.norm(xn, p, 1)
            self.assertEqual(res.shape, expected.shape)
            self.assertEqual(res, expected, msg="dim reduction failed for {}-norm".format(p))

        # matrix norm
        for p in ['fro', 'nuc']:
            res = x.norm(p).cpu()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res.shape, expected.shape)
            self.assertEqual(res, expected, msg="dim reduction failed for {}-norm".format(p))

        # larger tensor sanity check
        self.assertEqual(2 * torch.norm(torch.ones(10000)), torch.norm(torch.ones(40000)))
//  [python]
def test_norm(self, device):
        # full reduction
        x = torch.randn(25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = torch.linalg.norm(x, p).item()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res, expected, atol=1e-5, rtol=0, msg="full reduction failed for {}-norm".format(p))

        # one dimension
        x = torch.randn(25, 25, device=device)
        xn = x.cpu().numpy()
        for p in [0, 1, 2, 3, 4, inf, -inf]:
            res = torch.linalg.norm(x, p, 1).cpu()
            expected = np.linalg.norm(xn, p, 1)
            self.assertEqual(res.shape, expected.shape)
            self.assertEqual(res, expected, msg="dim reduction failed for {}-norm".format(p))

        # matrix norm
        for p in ['fro', 'nuc']:
            res = torch.linalg.norm(x, p).cpu()
            expected = np.linalg.norm(xn, p)
            self.assertEqual(res.shape, expected.shape)
            self.assertEqual(res, expected, msg="dim reduction failed for {}-norm".format(p))

        # larger tensor sanity check
        self.assertEqual(2 * torch.linalg.norm(torch.ones(10000)), torch.linalg.norm(torch.ones(40000)))
//  [python]
def loss_8c(self, normal):
    """
    normal: B x N x 3
    friction: B x N x 3
    x: B x N x 3
    mu: ()
    """
    normal = normal / torch.norm(normal, dim=-1, keepdim=True)
    left = torch.einsum('ijk, ijk->ij', normal, normal)
    right = torch.norm(normal, dim=2) / self.sqrt_sq_mu_1
    diff = left - right
    return torch.sum(self.relu(-diff), 1)
//  [python]
def loss_8c(self, normal):
    """
    normal: B x N x 3
    friction: B x N x 3
    x: B x N x 3
    mu: ()
    """
    normal = normal / torch.linalg.norm(normal, dim=-1, keepdim=True)
    left = torch.einsum('ijk, ijk->ij', normal, normal)
    right = torch.linalg.norm(normal, dim=2) / self.sqrt_sq_mu_1
    diff = left - right
    return torch.sum(self.relu(-diff), 1)
//  [python]
def depth2pts_outside(ray_o, ray_d, depth):
    '''
    ray_o, ray_d: [..., 3]
    depth: [...]; inverse of distance to sphere origin
    '''
    # note: d1 becomes negative if this mid point is behind camera
    d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
    p_mid = ray_o + d1.unsqueeze(-1) * ray_d
    p_mid_norm = torch.norm(p_mid, dim=-1)
    ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
    d2 = torch.sqrt(1. - p_mid_norm * p_mid_norm) * ray_d_cos
    p_sphere = ray_o + (d1 + d2).unsqueeze(-1) * ray_d

    rot_axis = torch.cross(ray_o, p_sphere, dim=-1)
    rot_axis = rot_axis / torch.norm(rot_axis, dim=-1, keepdim=True)
    phi = torch.asin(p_mid_norm)
    theta = torch.asin(p_mid_norm * depth)  # depth is inside [0, 1]
    rot_angle = (phi - theta).unsqueeze(-1)     # [..., 1]

    # now rotate p_sphere
    # Rodrigues formula: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
    p_sphere_new = p_sphere * torch.cos(rot_angle) + \
                torch.cross(rot_axis, p_sphere, dim=-1) * torch.sin(rot_angle) + \
                rot_axis * torch.sum(rot_axis*p_sphere, dim=-1, keepdim=True) * (1.-torch.cos(rot_angle))
    p_sphere_new = p_sphere_new / torch.norm(p_sphere_new, dim=-1, keepdim=True)
    pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)

    # now calculate conventional depth
    depth_real = 1. / (depth + TINY_NUMBER) * torch.cos(theta) * ray_d_cos + d1
    return pts, depth_real