jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def signguard(gradients, net, lr, f, byz, device, seed):
    """
    Based on the description in https://arxiv.org/abs/2109.05872
    gradients: list of gradients.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    seed: seed for randomness
    """
    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    if byz == attacks.fltrust_attack:
        param_list = byz(param_list, net, lr, f, device)[:-1]
    else:
        param_list = byz(param_list, net, lr, f, device)
    n = len(param_list)

    num_params = param_list[0].size(0)
    selection_fraction = 0.1

    # lower and upper bound L,R for gradient norm
    L = 0.1
    R = 3.0
    S1 = []
    S2 = []

    # compute l2-norm
    l2_norm = torch.stack([torch.linalg.norm(g.flatten(), ord=2) for g in param_list])

    # compute element wise sign
    num_selection = int(num_params * selection_fraction)
    perm = torch.randperm(num_params)
    idx = perm[:num_selection]
    sign_grads = [torch.sign(g[idx]) for g in param_list]

    # norm-threshold filtering
    M = torch.median(l2_norm)
    for i in range(n):
        if L <= l2_norm[i] / M and l2_norm[i] / M <= R:
            S1.append(i)

    # compute sign statistics
    sign_pos = torch.stack([grad.eq(1.0).float().mean() for grad in sign_grads])
    sign_zero = torch.stack([grad.eq(0.0).float().mean() for grad in sign_grads])
    sign_neg = torch.stack([grad.eq(-1.0).float().mean() for grad in sign_grads])

    # sign-based clustering
    sign_feat = torch.stack([sign_pos, sign_zero, sign_neg], dim=1).detach().cpu().numpy()
    cluster = KMeans(n_clusters=2, max_iter=20, random_state=seed)
    labels = cluster.fit_predict(sign_feat)

    labels_tensor = torch.from_numpy(labels).to(device)
    count = torch.bincount(labels_tensor)
    largest_cluster = torch.argmax(count)

    for i, value in enumerate(labels_tensor):
        if value == largest_cluster:
            S2.append(i)

    # compute intersection of S1 and S2
    S = [i for i in S1 if i in S2]

    # global update
    global_update = torch.zeros(param_list[0].size()).to(device)
    for i in S:
        global_update += param_list[i]
    global_update *= 1 / len(S)

    # update the global model
    idx = 0
    for j, param in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-lr)
        idx += torch.numel(param)
//  [python]
def __init__(self, u, cos_angle=.05):
        batch_size = u.size(0)
        # normalize the cone directions
        self.directions = utils.bmul(u, 1. / torch.norm(u.reshape(batch_size, -1), dim=-1))
        self.cos_angle = cos_angle
        self.alpha = np.sqrt(1. / cos_angle - 1)
//  [python]
def __init__(self, u, cos_angle=.05):
        batch_size = u.size(0)
        # normalize the cone directions
        self.directions = utils.bmul(u, 1. / torch.linalg.norm(u.reshape(batch_size, -1), dim=-1))
        self.cos_angle = cos_angle
        self.alpha = np.sqrt(1. / cos_angle - 1)
//  [python]
def lmo(self, grad, iterate):
        """Linear Maximization Oracle.
        Return s - iterate with s solving the linear problem

        ..math::
            max_{||s||_2 <= alpha} <grad, s>

        Args:
          grad: torch.Tensor of shape (batch_size, *)
              usually -gradient
          iterate: torch.Tensor of shape (batch_size, *)
              usually the iterate of the considered algorithm

        Returns:
          update_direction: torch.Tensor, same shape as grad and iterate,
              s - iterate, where s is the vertex of the constraint most correlated
              with u
          max_step_size: torch.Tensor of shape (batch_size,)
              1. for a Frank-Wolfe step.
        """
        update_direction = -iterate.clone().detach()
        grad_norms = torch.norm(grad.view(grad.size(0), -1), p=2, dim=-1)
        update_direction += self.alpha * (grad.view(grad.size(0), -1).T
                                            / grad_norms).T.view_as(iterate)
        return update_direction, torch.ones(iterate.size(0), device=iterate.device, dtype=iterate.dtype)
//  [python]
def lmo(self, grad, iterate):
        """Linear Maximization Oracle.
        Return s - iterate with s solving the linear problem

        ..math::
            max_{||s||_2 <= alpha} <grad, s>

        Args:
          grad: torch.Tensor of shape (batch_size, *)
              usually -gradient
          iterate: torch.Tensor of shape (batch_size, *)
              usually the iterate of the considered algorithm

        Returns:
          update_direction: torch.Tensor, same shape as grad and iterate,
              s - iterate, where s is the vertex of the constraint most correlated
              with u
          max_step_size: torch.Tensor of shape (batch_size,)
              1. for a Frank-Wolfe step.
        """
        update_direction = -iterate.clone().detach()
        grad_norms = torch.linalg.norm(grad.view(grad.size(0), -1), ord=2, dim=-1)
        update_direction += self.alpha * (grad.view(grad.size(0), -1).T
                                            / grad_norms).T.view_as(iterate)
        return update_direction, torch.ones(iterate.size(0), device=iterate.device, dtype=iterate.dtype)
//  [python]
def normalize_vector(v):  # v: [B, 3]
    batch = v.shape[0]
    v_mag = torch.norm(v, p=2, dim=1)  # batch
    """
    if torch.any(torch.isnan(v_mag)):
        print('nan in v_mag!')
        idx = torch.where(torch.isnan(v_mag))[0]
        print('v_mag', v_mag[idx], 'v', v[idx])
    if torch.any(torch.isinf(v_mag)):
        print('inf in v_mag!')
        idx = torch.where(torch.isinf(v_mag))[0]
        print('v_mag', v_mag[idx], 'v', v[idx])
    """
    eps = torch.autograd.Variable(torch.FloatTensor([1e-8]).to(v_mag.device))
    valid_mask = (v_mag > eps).float().view(batch, 1)
    backup = torch.tensor([1.0, 0.0, 0.0]).float().to(v.device).view(1, 3).expand(batch, 3)
    v_mag = torch.max(v_mag, eps)
    v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
    v = v / v_mag
    ret = v * valid_mask + backup * (1 - valid_mask)
    """
    if torch.any(torch.isnan(v)):
        print('nan in v!')
        idx = torch.where(torch.isnan(v))[0]
        print('v_mag', v_mag[idx], 'v', v[idx], 'valid', valid_mask[idx])
    if torch.any(torch.isnan(ret)):
        print('nan in ret!')
        idx = torch.where(torch.isnan(ret))[0]
        print('v_mag', v_mag[idx], 'v', v[idx], 'valid', valid_mask[idx])
    """
    return ret
//  [python]
def normalize_vector(v):  # v: [B, 3]
    batch = v.shape[0]
    v_mag = torch.linalg.norm(v, ord=2, dim=1)  # batch
    """
    if torch.any(torch.isnan(v_mag)):
        print('nan in v_mag!')
        idx = torch.where(torch.isnan(v_mag))[0]
        print('v_mag', v_mag[idx], 'v', v[idx])
    if torch.any(torch.isinf(v_mag)):
        print('inf in v_mag!')
        idx = torch.where(torch.isinf(v_mag))[0]
        print('v_mag', v_mag[idx], 'v', v[idx])
    """
    eps = torch.autograd.Variable(torch.FloatTensor([1e-8]).to(v_mag.device))
    valid_mask = (v_mag > eps).float().view(batch, 1)
    backup = torch.tensor([1.0, 0.0, 0.0]).float().to(v.device).view(1, 3).expand(batch, 3)
    v_mag = torch.max(v_mag, eps)
    v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
    v = v / v_mag
    ret = v * valid_mask + backup * (1 - valid_mask)
    """
    if torch.any(torch.isnan(v)):
        print('nan in v!')
        idx = torch.where(torch.isnan(v))[0]
        print('v_mag', v_mag[idx], 'v', v[idx], 'valid', valid_mask[idx])
    if torch.any(torch.isnan(ret)):
        print('nan in ret!')
        idx = torch.where(torch.isnan(ret))[0]
        print('v_mag', v_mag[idx], 'v', v[idx], 'valid', valid_mask[idx])
    """
    return ret
//  [python]
def test_frobenius_norm_with_dim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, dim=0)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_frobenius_norm_with_dim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, dim=0)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_frobenius_norm(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input).view(1)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        print(tfl_output.shape)
        print(dummy_output, tfl_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_frobenius_norm(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input).view(1)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        print(tfl_output.shape)
        print(dummy_output, tfl_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_frobenius_norm_with_dim_keepdim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, dim=0, keepdim=True)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_frobenius_norm_with_dim_keepdim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, dim=0, keepdim=True)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p1(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, p=1)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p1(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, ord=1)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p1_with_dim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, p=1, dim=0)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p1_with_dim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, ord=1, dim=0)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def prox(self, x, step_size=None):
        """
        Projects `x` batch-wise onto the cone constraint.
        Args:
          x: torch.Tensor of shape (batch_size, *)
            batch of vectors to project
          step_size: Any
            Not used
        Returns:
          proj_x: torch.Tensor of shape (batch_size, *)
            batch-wise projection of `x` onto the cone constraint.
        """
        batch_size = x.size(0)
        uTx = utils.bdot(self.directions, x)
        p_u = self.proj_u(x)
        p_orth_u = x - p_u
        norm_p_orth_u = torch.norm(p_orth_u.reshape(batch_size, -1), dim=-1)
        identity_idx = (norm_p_orth_u <= self.alpha * uTx)
        zero_idx = (self.alpha * norm_p_orth_u <= - uTx)
        project_idx = ~torch.logical_or(identity_idx, zero_idx)

        res = x.detach().clone()
        res[zero_idx] = 0.
        res[project_idx] = utils.bmul((self.alpha * norm_p_orth_u[project_idx] + uTx[project_idx]) / (1. + self.alpha ** 2), 
                                      (self.alpha * utils.bmul(p_orth_u[project_idx], 1 / norm_p_orth_u[project_idx])
                                       + self.directions[project_idx]))
        return res
//  [python]
def prox(self, x, step_size=None):
        """
        Projects `x` batch-wise onto the cone constraint.
        Args:
          x: torch.Tensor of shape (batch_size, *)
            batch of vectors to project
          step_size: Any
            Not used
        Returns:
          proj_x: torch.Tensor of shape (batch_size, *)
            batch-wise projection of `x` onto the cone constraint.
        """
        batch_size = x.size(0)
        uTx = utils.bdot(self.directions, x)
        p_u = self.proj_u(x)
        p_orth_u = x - p_u
        norm_p_orth_u = torch.linalg.norm(p_orth_u.reshape(batch_size, -1), dim=-1)
        identity_idx = (norm_p_orth_u <= self.alpha * uTx)
        zero_idx = (self.alpha * norm_p_orth_u <= - uTx)
        project_idx = ~torch.logical_or(identity_idx, zero_idx)

        res = x.detach().clone()
        res[zero_idx] = 0.
        res[project_idx] = utils.bmul((self.alpha * norm_p_orth_u[project_idx] + uTx[project_idx]) / (1. + self.alpha ** 2), 
                                      (self.alpha * utils.bmul(p_orth_u[project_idx], 1 / norm_p_orth_u[project_idx])
                                       + self.directions[project_idx]))
        return res
//  [python]
def test_norm_p1_with_dim_keepdim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, p=1, dim=0, keepdim=True)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p1_with_dim_keepdim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, ord=1, dim=0, keepdim=True)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p2(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, p=2)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input).view(1)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p2(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, ord=2)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input).view(1)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p2_with_dim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, p=2, dim=0)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p2_with_dim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, ord=2, dim=0)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p2_with_dim_keepdim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.norm(x, p=2, dim=0, keepdim=True)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def test_norm_p2_with_dim_keepdim(self):
        dummy_input = torch.randn(10, 10, dtype=torch.float32)

        def model(x):
            return torch.linalg.norm(x, ord=2, dim=0, keepdim=True)

        model_path = get_model_path()

        converter = TFLiteConverter(model, dummy_input, model_path, nchw_transpose=False)
        converter.convert()

        dummy_output = model(dummy_input)
        tfl_output = tfl_run_model(model_path, dummy_input, dummy_output)
        assert_close(dummy_output, tfl_output, atol=256.0, rtol=256.0)
//  [python]
def print_norm(x):
    print("printing norm 2", torch.norm(x))
//  [python]
def print_norm(x):
    print("printing norm 2", torch.linalg.norm(x))
//  [python]
def mpjpe(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def mpjpe(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def slerp(val, low, high):
    """
    Find the interpolation point between the 'low' and 'high' values for the given 'val'. See https://en.wikipedia.org/wiki/Slerp for more details on the topic.
    """
    low_norm = low / torch.norm(low)
    high_norm = high / torch.norm(high)
    omega = torch.acos((low_norm * high_norm))
    so = torch.sin(omega)
    res = (torch.sin((1.0 - val) * omega) / so) * low + (torch.sin(val * omega) / so) * high
    return res
//  [python]
def slerp(val, low, high):
    """
    Find the interpolation point between the 'low' and 'high' values for the given 'val'. See https://en.wikipedia.org/wiki/Slerp for more details on the topic.
    """
    low_norm = low / torch.linalg.norm(low)
    high_norm = high / torch.linalg.norm(high)
    omega = torch.acos((low_norm * high_norm))
    so = torch.sin(omega)
    res = (torch.sin((1.0 - val) * omega) / so) * low + (torch.sin(val * omega) / so) * high
    return res
//  [python]
def pairwise_dist(self, x, y):
        """For n points x_i and m points y_j computed dist(x_i,y_j)
            TODO: CHECK
        """
        x_y_ = self.pairwise_embed(x, y)
        log_x_y_ = torch.arccos(x_y_.real)
        dist = math.sqrt(2)*torch.minimum(log_x_y_, 2*pi-log_x_y_)
        dist = torch.norm(dist, dim=1).reshape(x.shape[0], y.shape[0])
        return dist
//  [python]
def pairwise_dist(self, x, y):
        """For n points x_i and m points y_j computed dist(x_i,y_j)
            TODO: CHECK
        """
        x_y_ = self.pairwise_embed(x, y)
        log_x_y_ = torch.arccos(x_y_.real)
        dist = math.sqrt(2)*torch.minimum(log_x_y_, 2*pi-log_x_y_)
        dist = torch.linalg.norm(dist, dim=1).reshape(x.shape[0], y.shape[0])
        return dist
//  [python]
def __call__(self, source, target, init_A=None):
        """
        :param source: Shape with points BxNxD
        :param target_batch: Shape with points BxMxD
        :return: Bx(D+1)xD transform matrix
        """
        source, target = self.sampling_input(source, target)
        toflow = source
        A_prev = init_A if init_A is not None else None
        A = None
        if self.search_init_transform:
            A_prev, toflow = self.find_initial_transform(source, target)
        for i in range(self.niter):
            toflow, target = self.extract_fea(toflow, target, i)
            flowed, weight_map_ratio = self.get_correspondence_shape(toflow, target)
            if not self.use_barycenter_weight:
                A, transforme_points = self._solve_transform(toflow, flowed)
            else:
                toflow_weights = toflow.weights
                toflow.weights = weight_map_ratio
                A, transforme_points = self._solve_transform(toflow, flowed)
                toflow.weights = toflow_weights
            A = self.compose_transform(A_prev, A) if A_prev is not None else A
            transformed_points = (
                torch.cat(
                    (source.points, torch.ones_like(source.points[:, :, :1])), dim=2
                )
                @ A
            )
            toflow = Shape().set_data_with_refer_to(transformed_points, source)
            if i > 0 and torch.norm(A - A_prev) < self.rel_ftol:
                print(
                    "reach relative tolerance {}".format(torch.norm(A - A_prev).item())
                )
                break
            A_prev = A
            if self.plot:
                self.visualize(
                    source, toflow, target, weight_map_ratio, self.geomloss_setting, i
                )
        return A
//  [python]
def __call__(self, source, target, init_A=None):
        """
        :param source: Shape with points BxNxD
        :param target_batch: Shape with points BxMxD
        :return: Bx(D+1)xD transform matrix
        """
        source, target = self.sampling_input(source, target)
        toflow = source
        A_prev = init_A if init_A is not None else None
        A = None
        if self.search_init_transform:
            A_prev, toflow = self.find_initial_transform(source, target)
        for i in range(self.niter):
            toflow, target = self.extract_fea(toflow, target, i)
            flowed, weight_map_ratio = self.get_correspondence_shape(toflow, target)
            if not self.use_barycenter_weight:
                A, transforme_points = self._solve_transform(toflow, flowed)
            else:
                toflow_weights = toflow.weights
                toflow.weights = weight_map_ratio
                A, transforme_points = self._solve_transform(toflow, flowed)
                toflow.weights = toflow_weights
            A = self.compose_transform(A_prev, A) if A_prev is not None else A
            transformed_points = (
                torch.cat(
                    (source.points, torch.ones_like(source.points[:, :, :1])), dim=2
                )
                @ A
            )
            toflow = Shape().set_data_with_refer_to(transformed_points, source)
            if i > 0 and torch.linalg.norm(A - A_prev) < self.rel_ftol:
                print(
                    "reach relative tolerance {}".format(torch.linalg.norm(A - A_prev).item())
                )
                break
            A_prev = A
            if self.plot:
                self.visualize(
                    source, toflow, target, weight_map_ratio, self.geomloss_setting, i
                )
        return A
//  [python]
def _update_metric(self):
        # convert to 1 big tensor

        self.M_tens = torch.cat(list(self.M))
        self.centroids_tens = torch.cat(list(self.centroids))

        # define new metric
        def G(z):
            return torch.inverse(
                (
                    self.M_tens.unsqueeze(0)
                    * torch.exp(
                        -torch.norm(
                            self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                        )
                        ** 2
                        / (self.temperature**2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1)
                + self.lbd * torch.eye(self.latent_dim).to(z.device)
            )

        def G_inv(z):
            return (
                self.M_tens.unsqueeze(0)
                * torch.exp(
                    -torch.norm(
                        self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                    )
                    ** 2
                    / (self.temperature**2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(z.device)

        self.G = G
        self.G_inv = G_inv
        self.M = deque(maxlen=100)
        self.centroids = deque(maxlen=100)
//  [python]
def _update_metric(self):
        # convert to 1 big tensor

        self.M_tens = torch.cat(list(self.M))
        self.centroids_tens = torch.cat(list(self.centroids))

        # define new metric
        def G(z):
            return torch.inverse(
                (
                    self.M_tens.unsqueeze(0)
                    * torch.exp(
                        -torch.linalg.norm(
                            self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                        )
                        ** 2
                        / (self.temperature**2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1)
                + self.lbd * torch.eye(self.latent_dim).to(z.device)
            )

        def G_inv(z):
            return (
                self.M_tens.unsqueeze(0)
                * torch.exp(
                    -torch.linalg.norm(
                        self.centroids_tens.unsqueeze(0) - z.unsqueeze(1), dim=-1
                    )
                    ** 2
                    / (self.temperature**2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(z.device)

        self.G = G
        self.G_inv = G_inv
        self.M = deque(maxlen=100)
        self.centroids = deque(maxlen=100)
//  [python]
def validate_rot_mat_torch(R, eps=1e-6):
    Rt = torch.transpose(R, 0, 1)
    expect_I = torch.matmul(Rt, R)
    I = torch.eye(3)
    n = torch.norm(I - expect_I)
    ortho_mat = n < eps

    pos_det = np.abs(torch.det(R) - 1) < eps

    return ortho_mat and pos_det
//  [python]
def validate_rot_mat_torch(R, eps=1e-6):
    Rt = torch.transpose(R, 0, 1)
    expect_I = torch.matmul(Rt, R)
    I = torch.eye(3)
    n = torch.linalg.norm(I - expect_I)
    ortho_mat = n < eps

    pos_det = np.abs(torch.det(R) - 1) < eps

    return ortho_mat and pos_det
//  [python]
def fit_cylinder_torch(self, points, normals, weights, ids=0, show_warning=False):
        # compute
        # U, s, V = torch.svd(weights * normals)
        weighted_normals = weights * normals

        if np.linalg.cond(weighted_normals.data.cpu().numpy()) > 1e5:
            if show_warning:
                print("condition number is large in cylinder")
                print(torch.sum(normals).item(), torch.sum(points).item(), torch.sum(weights).item())

        U, s, V = customsvd(weighted_normals)
        a = V[:, -1]
        a = torch.reshape(a, (1, 3))

        # find the projection onto a plane perpendicular to the axis
        a = a.reshape((3, 1))
        a = a / (torch.norm(a, 2) + EPS)

        prj_circle = points - ((points @ a).permute(1, 0) * a).permute(1, 0)

        # torch doesn't have least square for
        center, radius = self.fit_sphere_torch(prj_circle, normals, weights)
        return a, center, radius
//  [python]
def fit_cylinder_torch(self, points, normals, weights, ids=0, show_warning=False):
        # compute
        # U, s, V = torch.svd(weights * normals)
        weighted_normals = weights * normals

        if np.linalg.cond(weighted_normals.data.cpu().numpy()) > 1e5:
            if show_warning:
                print("condition number is large in cylinder")
                print(torch.sum(normals).item(), torch.sum(points).item(), torch.sum(weights).item())

        U, s, V = customsvd(weighted_normals)
        a = V[:, -1]
        a = torch.reshape(a, (1, 3))

        # find the projection onto a plane perpendicular to the axis
        a = a.reshape((3, 1))
        a = a / (torch.linalg.norm(a, 2) + EPS)

        prj_circle = points - ((points @ a).permute(1, 0) * a).permute(1, 0)

        # torch doesn't have least square for
        center, radius = self.fit_sphere_torch(prj_circle, normals, weights)
        return a, center, radius
//  [python]
def intersect_sphere(rays_o, rays_d):
    """Compute the depth of the intersection point between this ray and unit sphere.
    Args:
        rays_o: [num_rays, 3]. Ray origins.
        rays_d: [num_rays, 3]. Ray directions.
    Returns:
        depth: [num_rays, 1]. Depth of the intersection point.
    """
    # note: d1 becomes negative if this mid point is behind camera

    d1 = -torch.sum(rays_d * rays_o, dim=-1, keepdim=True) / torch.sum(
        rays_d**2, dim=-1, keepdim=True
    )
    p = rays_o + d1 * rays_d
    # consider the case where the ray does not intersect the sphere
    rays_d_cos = 1.0 / torch.norm(rays_d, dim=-1, keepdim=True)
    p_norm_sq = torch.sum(p * p, dim=-1, keepdim=True)
    d2 = torch.sqrt(1.0 - p_norm_sq) * rays_d_cos

    return d1 + d2
//  [python]
def intersect_sphere(rays_o, rays_d):
    """Compute the depth of the intersection point between this ray and unit sphere.
    Args:
        rays_o: [num_rays, 3]. Ray origins.
        rays_d: [num_rays, 3]. Ray directions.
    Returns:
        depth: [num_rays, 1]. Depth of the intersection point.
    """
    # note: d1 becomes negative if this mid point is behind camera

    d1 = -torch.sum(rays_d * rays_o, dim=-1, keepdim=True) / torch.sum(
        rays_d**2, dim=-1, keepdim=True
    )
    p = rays_o + d1 * rays_d
    # consider the case where the ray does not intersect the sphere
    rays_d_cos = 1.0 / torch.linalg.norm(rays_d, dim=-1, keepdim=True)
    p_norm_sq = torch.sum(p * p, dim=-1, keepdim=True)
    d2 = torch.sqrt(1.0 - p_norm_sq) * rays_d_cos

    return d1 + d2
//  [python]
def shrink_group(self, ts):
		# shrinkage for 4-d weight(conv layer)
		ts_reshape = torch.reshape(ts,(ts.shape[0],-1))
		ts_norm = torch.norm(ts_reshape,2,1)
		ts_shrink = torch.max(torch.zeros_like(ts_norm),torch.ones_like(ts_norm) - torch.div(torch.ones_like(ts_norm),ts_norm))
		ts_return = torch.transpose(torch.mul(torch.transpose(ts_reshape,0,1),ts_shrink),0,1)
		ts_return = torch.reshape(ts_return,ts.shape)
		return ts_return
//  [python]
def shrink_group(self, ts):
		# shrinkage for 4-d weight(conv layer)
		ts_reshape = torch.reshape(ts,(ts.shape[0],-1))
		ts_norm = torch.linalg.norm(ts_reshape,2,1)
		ts_shrink = torch.max(torch.zeros_like(ts_norm),torch.ones_like(ts_norm) - torch.div(torch.ones_like(ts_norm),ts_norm))
		ts_return = torch.transpose(torch.mul(torch.transpose(ts_reshape,0,1),ts_shrink),0,1)
		ts_return = torch.reshape(ts_return,ts.shape)
		return ts_return
//  [python]
def busemann(x, p, keepdim=True):
    """
    x: (..., d)
    p: (..., d)

    Returns: (..., 1) if keepdim==True else (...)
    """

    xnorm = x.norm(dim=-1, p=2, keepdim=True)
    pnorm = p.norm(dim=-1, p=2, keepdim=True)
    p = p / pnorm.clamp_min(MIN_NORM)
    num = torch.norm(p - x, dim=-1, keepdim=True) ** 2
    den = (1 - xnorm ** 2).clamp_min(MIN_NORM)
    ans = torch.log((num / den).clamp_min(MIN_NORM))
    if not keepdim:
        ans = ans.squeeze(-1)
    return ans
//  [python]
def busemann(x, p, keepdim=True):
    """
    x: (..., d)
    p: (..., d)

    Returns: (..., 1) if keepdim==True else (...)
    """

    xnorm = torch.linalg.norm(x, dim=-1, keepdim=True)
    pnorm = torch.linalg.norm(p, dim=-1, keepdim=True)
    p = p / pnorm.clamp_min(MIN_NORM)
    num = torch.linalg.norm(p - x, dim=-1, keepdim=True) ** 2
    den = (1 - xnorm ** 2).clamp_min(MIN_NORM)
    ans = torch.log((num / den).clamp_min(MIN_NORM))
    if not keepdim:
        ans = ans.squeeze(-1)
    return ans
//  [python]
def busemann_to_horocycle(p, t):
    """ Find the horocycle corresponding to the level set of the Busemann function to ideal point p with value t.

    p: (..., d)
    t: (...)

    Returns:
    c: (..., d)
    r: (...)
    """
    # Busemann_p(x) = d means dist(0, x) = -d
    q = -torch.tanh(t / 2).unsqueeze(-1) * p
    c = (p + q) / 2.0
    r = torch.norm(p - q, dim=-1) / 2.0
    return c, r
//  [python]
def busemann_to_horocycle(p, t):
    """ Find the horocycle corresponding to the level set of the Busemann function to ideal point p with value t.

    p: (..., d)
    t: (...)

    Returns:
    c: (..., d)
    r: (...)
    """
    # Busemann_p(x) = d means dist(0, x) = -d
    q = -torch.tanh(t / 2).unsqueeze(-1) * p
    c = (p + q) / 2.0
    r = torch.linalg.norm(p - q, dim=-1) / 2.0
    return c, r
//  [python]
def get_rays(directions, c2w, output_view_dirs = False, output_radii = False):
    """
    Get ray origin and normalized directions in world coordinate for all pixels in one image.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Inputs:
        directions: (H, W, 3) precomputed ray directions in camera coordinate
        c2w: (3, 4) transformation matrix from camera coordinate to world coordinate

    Outputs:
        rays_o: (H*W, 3), the origin of the rays in world coordinate
        rays_d: (H*W, 3), the normalized direction of the rays in world coordinate
    """
    # Rotate ray directions from camera coordinate to the world coordinate
    rays_d = directions @ c2w[:, :3].T # (H, W, 3)
    #rays_d /= torch.norm(rays_d, dim=-1, keepdim=True)
    # The origin of all rays is the camera origin in world coordinate
    rays_o = c2w[:, 3].expand(rays_d.shape) # (H, W, 3)

    if output_radii:
        rays_d_orig = directions @ c2w[:, :3].T
        dx = torch.sqrt(torch.sum((rays_d_orig[:-1, :, :] - rays_d_orig[1:, :, :]) ** 2, dim=-1))
        dx = torch.cat([dx, dx[-2:-1, :]], dim=0)
        radius = dx[..., None] * 2 / torch.sqrt(torch.tensor(12, dtype=torch.int8))
        radius = radius.reshape(-1)
    
    if output_view_dirs:
        viewdirs = rays_d
        viewdirs /= torch.norm(viewdirs, dim=-1, keepdim=True)
        rays_d = rays_d.view(-1, 3)
        rays_o = rays_o.view(-1, 3)
        viewdirs = viewdirs.view(-1, 3)
        if output_radii:
            return rays_o, viewdirs, rays_d, radius
        else:
            return rays_o, viewdirs, rays_d  
    else:
        rays_d /= torch.norm(rays_d, dim=-1, keepdim=True)
        rays_d = rays_d.view(-1, 3)
        rays_o = rays_o.view(-1, 3)
        return rays_o, rays_d
//  [python]
def get_rays(directions, c2w, output_view_dirs = False, output_radii = False):
    """
    Get ray origin and normalized directions in world coordinate for all pixels in one image.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Inputs:
        directions: (H, W, 3) precomputed ray directions in camera coordinate
        c2w: (3, 4) transformation matrix from camera coordinate to world coordinate

    Outputs:
        rays_o: (H*W, 3), the origin of the rays in world coordinate
        rays_d: (H*W, 3), the normalized direction of the rays in world coordinate
    """
    # Rotate ray directions from camera coordinate to the world coordinate
    rays_d = directions @ c2w[:, :3].T # (H, W, 3)
    #rays_d /= torch.norm(rays_d, dim=-1, keepdim=True)
    # The origin of all rays is the camera origin in world coordinate
    rays_o = c2w[:, 3].expand(rays_d.shape) # (H, W, 3)

    if output_radii:
        rays_d_orig = directions @ c2w[:, :3].T
        dx = torch.sqrt(torch.sum((rays_d_orig[:-1, :, :] - rays_d_orig[1:, :, :]) ** 2, dim=-1))
        dx = torch.cat([dx, dx[-2:-1, :]], dim=0)
        radius = dx[..., None] * 2 / torch.sqrt(torch.tensor(12, dtype=torch.int8))
        radius = radius.reshape(-1)
    
    if output_view_dirs:
        viewdirs = rays_d
        viewdirs /= torch.linalg.norm(viewdirs, dim=-1, keepdim=True)
        rays_d = rays_d.view(-1, 3)
        rays_o = rays_o.view(-1, 3)
        viewdirs = viewdirs.view(-1, 3)
        if output_radii:
            return rays_o, viewdirs, rays_d, radius
        else:
            return rays_o, viewdirs, rays_d  
    else:
        rays_d /= torch.linalg.norm(rays_d, dim=-1, keepdim=True)
        rays_d = rays_d.view(-1, 3)
        rays_o = rays_o.view(-1, 3)
        return rays_o, rays_d
//  [python]
def sphere_intersection(c1, r1, c2, r2):
    """ Computes the intersections of a circle centered at ci of radius ri.

    c1, c2: (..., d)
    r1, r2: (...)

    Returns:
    center, radius such that the intersection of the two spheres is given by
    the intersection of the sphere (c, r) with the hyperplane orthogonal to the direction c1->c2
    """

    d = torch.norm(c1 - c2, dim=-1)  # (...)
    x, h = circle_intersection_(r1 / d.clamp_min(MIN_NORM), r2 / d.clamp_min(MIN_NORM))  # (...)
    x = x.unsqueeze(-1)
    center = x * c2 + (1 - x) * c1  # (..., d)
    radius = h * d  # (...)
    return center, radius
//  [python]
def sphere_intersection(c1, r1, c2, r2):
    """ Computes the intersections of a circle centered at ci of radius ri.

    c1, c2: (..., d)
    r1, r2: (...)

    Returns:
    center, radius such that the intersection of the two spheres is given by
    the intersection of the sphere (c, r) with the hyperplane orthogonal to the direction c1->c2
    """

    d = torch.linalg.norm(c1 - c2, dim=-1)  # (...)
    x, h = circle_intersection_(r1 / d.clamp_min(MIN_NORM), r2 / d.clamp_min(MIN_NORM))  # (...)
    x = x.unsqueeze(-1)
    center = x * c2 + (1 - x) * c1  # (..., d)
    radius = h * d  # (...)
    return center, radius
//  [python]
def circle_intersection(c1, c2, r1, r2):
    """ Computes the intersections of a circle centered at ci of radius ri.

    c1, c2: (..., d)
    r1, r2: (...)
    """

    d = torch.norm(c1 - c2)  # (...)
    x, h = circle_intersection_(r1 / d.clamp_min(MIN_NORM), r2 / d.clamp_min(MIN_NORM))  # (...)
    x = x.unsqueeze(-1)
    h = h.unsqueeze(-1)
    center = x * c2 + (1 - x) * c1  # (..., d)
    radius = h * d  # (...)

    # The intersection is a hypersphere of one lower dimension, intersected with the plane
    # orthogonal to the direction c1->c2
    # In general, you can compute this with a sort of higher dimensional cross product?
    # For now, only 2 dimensions

    ortho = c2 - c1  # (..., d)
    assert ortho.size(-1) == 2
    direction = torch.stack((-ortho[..., 1], ortho[..., 0]), dim=-1)
    direction = direction / torch.norm(direction, keepdim=True).clamp_min(MIN_NORM)
    return center + radius.unsqueeze(-1) * direction
//  [python]
def circle_intersection(c1, c2, r1, r2):
    """ Computes the intersections of a circle centered at ci of radius ri.

    c1, c2: (..., d)
    r1, r2: (...)
    """

    d = torch.linalg.norm(c1 - c2)  # (...)
    x, h = circle_intersection_(r1 / d.clamp_min(MIN_NORM), r2 / d.clamp_min(MIN_NORM))  # (...)
    x = x.unsqueeze(-1)
    h = h.unsqueeze(-1)
    center = x * c2 + (1 - x) * c1  # (..., d)
    radius = h * d  # (...)

    # The intersection is a hypersphere of one lower dimension, intersected with the plane
    # orthogonal to the direction c1->c2
    # In general, you can compute this with a sort of higher dimensional cross product?
    # For now, only 2 dimensions

    ortho = c2 - c1  # (..., d)
    assert ortho.size(-1) == 2
    direction = torch.stack((-ortho[..., 1], ortho[..., 0]), dim=-1)
    direction = direction / torch.linalg.norm(direction, keepdim=True).clamp_min(MIN_NORM)
    return center + radius.unsqueeze(-1) * direction
//  [python]
def perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_nocond, neg_scale, cond_scale):
    pos = noise_pred_pos - noise_pred_nocond
    neg = noise_pred_neg - noise_pred_nocond

    perp = neg - ((torch.mul(neg, pos).sum())/(torch.norm(pos)**2)) * pos
    perp_neg = perp * neg_scale
    cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg)
    return cfg_result
//  [python]
def perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_nocond, neg_scale, cond_scale):
    pos = noise_pred_pos - noise_pred_nocond
    neg = noise_pred_neg - noise_pred_nocond

    perp = neg - ((torch.mul(neg, pos).sum())/(torch.linalg.norm(pos)**2)) * pos
    perp_neg = perp * neg_scale
    cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg)
    return cfg_result
//  [python]
def depth2pts_outside(rays_o, rays_d, depth):
    """Compute the points along the ray that are outside of the unit sphere.
    Args:
        rays_o: [num_rays, 3]. Ray origins of the points.
        rays_d: [num_rays, 3]. Ray directions of the points.
        depth: [num_rays, num_samples along ray]. Inverse of distance to sphere origin.
    Returns:
        pts: [num_rays, 4]. Points outside of the unit sphere. (x', y', z', 1/r)
    """
    # note: d1 becomes negative if this mid point is behind camera
    rays_o = rays_o[..., None, :].expand(
        list(depth.shape) + [3]
    )  #  [N_rays, num_samples, 3]
    rays_d = rays_d[..., None, :].expand(
        list(depth.shape) + [3]
    )  #  [N_rays, num_samples, 3]
    d1 = -torch.sum(rays_d * rays_o, dim=-1, keepdim=True) / torch.sum(
        rays_d**2, dim=-1, keepdim=True
    )

    p_mid = rays_o + d1 * rays_d
    p_mid_norm = torch.norm(p_mid, dim=-1, keepdim=True)
    rays_d_cos = 1.0 / torch.norm(rays_d, dim=-1, keepdim=True)

    check_pos = 1.0 - p_mid_norm * p_mid_norm
    assert torch.all(check_pos >= 0), "1.0 - p_mid_norm * p_mid_norm should be greater than 0"

    d2 = torch.sqrt(1.0 - p_mid_norm * p_mid_norm) * rays_d_cos
    p_sphere = rays_o + (d1 + d2) * rays_d

    rot_axis = torch.cross(rays_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[..., None])  # depth is inside [0, 1]
    rot_angle = phi - theta  # [..., 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.0 - torch.cos(rot_angle))
    )
    p_sphere_new = p_sphere_new / (
        torch.norm(p_sphere_new, dim=-1, keepdim=True) + 1e-10
    )
    pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)

    return pts
//  [python]
def depth2pts_outside(rays_o, rays_d, depth):
    """Compute the points along the ray that are outside of the unit sphere.
    Args:
        rays_o: [num_rays, 3]. Ray origins of the points.
        rays_d: [num_rays, 3]. Ray directions of the points.
        depth: [num_rays, num_samples along ray]. Inverse of distance to sphere origin.
    Returns:
        pts: [num_rays, 4]. Points outside of the unit sphere. (x', y', z', 1/r)
    """
    # note: d1 becomes negative if this mid point is behind camera
    rays_o = rays_o[..., None, :].expand(
        list(depth.shape) + [3]
    )  #  [N_rays, num_samples, 3]
    rays_d = rays_d[..., None, :].expand(
        list(depth.shape) + [3]
    )  #  [N_rays, num_samples, 3]
    d1 = -torch.sum(rays_d * rays_o, dim=-1, keepdim=True) / torch.sum(
        rays_d**2, dim=-1, keepdim=True
    )

    p_mid = rays_o + d1 * rays_d
    p_mid_norm = torch.linalg.norm(p_mid, dim=-1, keepdim=True)
    rays_d_cos = 1.0 / torch.linalg.norm(rays_d, dim=-1, keepdim=True)

    check_pos = 1.0 - p_mid_norm * p_mid_norm
    assert torch.all(check_pos >= 0), "1.0 - p_mid_norm * p_mid_norm should be greater than 0"

    d2 = torch.sqrt(1.0 - p_mid_norm * p_mid_norm) * rays_d_cos
    p_sphere = rays_o + (d1 + d2) * rays_d

    rot_axis = torch.cross(rays_o, p_sphere, dim=-1)
    rot_axis = rot_axis / torch.linalg.norm(rot_axis, dim=-1, keepdim=True)
    phi = torch.asin(p_mid_norm)
    theta = torch.asin(p_mid_norm * depth[..., None])  # depth is inside [0, 1]
    rot_angle = phi - theta  # [..., 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.0 - torch.cos(rot_angle))
    )
    p_sphere_new = p_sphere_new / (
        torch.linalg.norm(p_sphere_new, dim=-1, keepdim=True) + 1e-10
    )
    pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)

    return pts