jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def weighted_svd_decomposition(
    A: torch.Tensor,
    W: torch.Tensor,
    heuristic: Optional[str],
    randomized: bool,
    num_ranks: int,
    num_oversampling: int = 5,
    normalize: bool = False,
    reduce_before_sqrt: bool = True,  # seems to be better empirically
) -> Tuple[torch.Tensor, torch.Tensor]:
    if A.ndim != 2 or W.ndim != 2:
        raise ValueError(f"Expected 2D Matrix, but got {A.ndim} and {W.ndim}.")
    if A.shape != W.shape:
        raise ValueError(f"Expected A.shape == W.shape, but got {A.shape} and {W.shape}.")

    if heuristic is None:
        heuristic = "two-sided"
    if heuristic not in ["none", "row", "col", "two-sided"]:
        raise ValueError(f"Heuristic {heuristic} not supported.")

    if normalize is True:
        W = W / torch.linalg.norm(W, ord="fro")

    if reduce_before_sqrt is True:
        # (A.shape[0], 1)
        W1 = torch.sqrt(torch.mean(W, dim=1, keepdim=True))
        # (1, A.shape[1])
        W2 = torch.sqrt(torch.mean(W, dim=0, keepdim=True))
    else:
        # (A.shape[0], 1)
        W1 = torch.mean(torch.sqrt(W), dim=1, keepdim=True)
        # (1, A.shape[1])
        W2 = torch.mean(torch.sqrt(W), dim=0, keepdim=True)

    if heuristic == "none":
        A_tilde = A
    elif heuristic == "row":
        A_tilde = W1 * A
    elif heuristic == "col":
        A_tilde = A * W2
    elif heuristic == "two-sided":
        A_tilde = W1 * A * W2
    else:
        raise ValueError

    L1_tilde, L2_tilde = svd_decomposition(
        A_tilde,
        randomized=randomized,
        num_ranks=num_ranks,
        num_oversampling=num_oversampling)

    if heuristic == "none":
        L1 = L1_tilde
        L2 = L2_tilde
    elif heuristic == "row":
        L1 = L1_tilde / W1
        L2 = L2_tilde
    elif heuristic == "col":
        L1 = L1_tilde
        L2 = L2_tilde / W2
    elif heuristic == "two-sided":
        L1 = L1_tilde / W1
        L2 = L2_tilde / W2
    else:
        raise ValueError

    return L1, L2
//  [python]
def min_angle_to_main_direction(main_dir,seg,nsample=32):
    # main_dir: 256,3,3, seg:256,64,2,3
    seg = seg[:,:,0,:] - seg[:,:,1,:] # 256,64,3
    seg_norm = seg / torch.norm(seg,dim=-1,keepdim=True) #256,64,3
    # seg_norm = seg_norm.unsqueeze(2) # 256,64,1,3

    dot = main_dir[:,None,:,:] * seg_norm[:,:,None,:] #256,64,3,3
    dot = torch.sum(dot,dim=-1) # 256,64,3
    dot = torch.abs(dot) # no direction angle
    e = 1e-8
    dot = torch.clip(dot,min=e,max=1-e)
    dot = torch.arccos(dot) # 256,64,3
    dot = torch.min(dot,dim=-1)[0] # min of three main direction 256,64
    return dot
//  [python]
def min_angle_to_main_direction(main_dir,seg,nsample=32):
    # main_dir: 256,3,3, seg:256,64,2,3
    seg = seg[:,:,0,:] - seg[:,:,1,:] # 256,64,3
    seg_norm = seg / torch.linalg.norm(seg,dim=-1,keepdim=True) #256,64,3
    # seg_norm = seg_norm.unsqueeze(2) # 256,64,1,3

    dot = main_dir[:,None,:,:] * seg_norm[:,:,None,:] #256,64,3,3
    dot = torch.sum(dot,dim=-1) # 256,64,3
    dot = torch.abs(dot) # no direction angle
    e = 1e-8
    dot = torch.clip(dot,min=e,max=1-e)
    dot = torch.arccos(dot) # 256,64,3
    dot = torch.min(dot,dim=-1)[0] # min of three main direction 256,64
    return dot
//  [python]
def pointLineSegDis(xyz,xyzSeg):
    # 点到线段的距离
    _,S,_ = xyz.shape # [1,S,3], S=256
    _,N,_,_ = xyzSeg.shape # [1,N,2,3]
    A = xyzSeg[:,:,0,:] # [1,N,3]
    B = xyzSeg[:,:,1,:]
    P = xyz # [1,S,3]
    AB = B - A # [1,N,3]
    AB_norm = AB / (torch.norm(AB,dim=-1,keepdim=True)+1e-8) #[1,N,3]

    PA = A[:,None,:,:] - P[:,:,None,:]  # [1,S,N,3]
    BP = P[:,:,None,:] - B[:,None,:,:]

    s = torch.multiply(PA,AB_norm.unsqueeze(1)).sum(-1)
    t = torch.multiply(BP, AB_norm.unsqueeze(1)).sum(-1)
    z = torch.zeros_like(s).to(P.device).float()

    tmp = torch.stack([s, t, z], dim=-1)
    h = torch.max(tmp,dim=-1)[0]

    AP = P[:,:,None,:] - A[:,None,:,:]

    AB_norm_repeat = AB_norm.unsqueeze(1).repeat(1,S,1,1)
    cross_result = torch.cross(AB_norm_repeat,AP,dim=-1)
    c = torch.norm(cross_result,dim=-1)
    return torch.sqrt(c**2+h**2)
//  [python]
def pointLineSegDis(xyz,xyzSeg):
    # 点到线段的距离
    _,S,_ = xyz.shape # [1,S,3], S=256
    _,N,_,_ = xyzSeg.shape # [1,N,2,3]
    A = xyzSeg[:,:,0,:] # [1,N,3]
    B = xyzSeg[:,:,1,:]
    P = xyz # [1,S,3]
    AB = B - A # [1,N,3]
    AB_norm = AB / (torch.linalg.norm(AB,dim=-1,keepdim=True)+1e-8) #[1,N,3]

    PA = A[:,None,:,:] - P[:,:,None,:]  # [1,S,N,3]
    BP = P[:,:,None,:] - B[:,None,:,:]

    s = torch.multiply(PA,AB_norm.unsqueeze(1)).sum(-1)
    t = torch.multiply(BP, AB_norm.unsqueeze(1)).sum(-1)
    z = torch.zeros_like(s).to(P.device).float()

    tmp = torch.stack([s, t, z], dim=-1)
    h = torch.max(tmp,dim=-1)[0]

    AP = P[:,:,None,:] - A[:,None,:,:]

    AB_norm_repeat = AB_norm.unsqueeze(1).repeat(1,S,1,1)
    cross_result = torch.cross(AB_norm_repeat,AP,dim=-1)
    c = torch.linalg.norm(cross_result,dim=-1)
    return torch.sqrt(c**2+h**2)
//  [python]
def normalize_embeds(*embeds, backend):
    normed_embeds = []
    for e in embeds:
        if backend == "tf":
            ne = tf.linalg.l2_normalize(e, axis=1)
        elif backend == "torch":
            norms = torch.norm(e, dim=1, keepdim=True)
            ne = e / norms
        else:
            norms = np.linalg.norm(e, axis=1, keepdims=True)
            ne = e / norms
        normed_embeds.append(ne)
    return normed_embeds[0] if len(embeds) == 1 else normed_embeds
//  [python]
def normalize_embeds(*embeds, backend):
    normed_embeds = []
    for e in embeds:
        if backend == "tf":
            ne = tf.linalg.l2_normalize(e, axis=1)
        elif backend == "torch":
            norms = torch.linalg.norm(e, dim=1, keepdim=True)
            ne = e / norms
        else:
            norms = np.linalg.norm(e, axis=1, keepdims=True)
            ne = e / norms
        normed_embeds.append(ne)
    return normed_embeds[0] if len(embeds) == 1 else normed_embeds
//  [python]
def angle_line_to_line(X,Y):
    # X: 1,N,2,3
    # Y: 1,S,2,3
    # 计算输入的线段pairwise的夹角， 返回余弦距离
    X,Y = X[0],Y[0]
    X_dir = X[:,0] - X[:, 1] # N,3
    Y_dir = Y[:,0] - Y[:, 1] # S,3
    X_dir_norm = X_dir/torch.norm(X_dir,dim=-1,keepdim=True)
    Y_dir_norm = Y_dir / torch.norm(Y_dir, dim=-1, keepdim=True)
    dot = X_dir_norm[:,None,:] * Y_dir_norm[None,:,:]
    dot = torch.sum(dot,dim=-1) # N,S
    dot = torch.abs(dot)
    e = 1e-7
    dot = torch.clip(dot,min=e,max=1-e)
    angle = torch.arccos(dot)
    angle = torch.rad2deg(angle)
    return angle.unsqueeze(0)
//  [python]
def angle_line_to_line(X,Y):
    # X: 1,N,2,3
    # Y: 1,S,2,3
    # 计算输入的线段pairwise的夹角， 返回余弦距离
    X,Y = X[0],Y[0]
    X_dir = X[:,0] - X[:, 1] # N,3
    Y_dir = Y[:,0] - Y[:, 1] # S,3
    X_dir_norm = X_dir/torch.linalg.norm(X_dir,dim=-1,keepdim=True)
    Y_dir_norm = Y_dir / torch.linalg.norm(Y_dir, dim=-1, keepdim=True)
    dot = X_dir_norm[:,None,:] * Y_dir_norm[None,:,:]
    dot = torch.sum(dot,dim=-1) # N,S
    dot = torch.abs(dot)
    e = 1e-7
    dot = torch.clip(dot,min=e,max=1-e)
    angle = torch.arccos(dot)
    angle = torch.rad2deg(angle)
    return angle.unsqueeze(0)
//  [python]
def farthest_point_sampleV4(xyz, npoint):
    """
    这个应该是基于角度进行FPS
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
        oldidx: [1,M]
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    def dist(center,other):
        pass
    seg = xyz.view(1,-1,2,3)
    device = xyz.device
    B, N, _, C = seg.shape
    if npoint>=N:
        x = np.arange(npoint)%N
        return torch.from_numpy(x).unsqueeze(0)
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device) * 0
    batch_indices = torch.arange(B, dtype=torch.long).to(device)

    seg_dir = seg[:,:,0,:] - seg[:,:,1,:]
    seg_dir_norm = seg_dir / torch.norm(seg_dir,dim=-1,keepdim=True)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid_line = seg[batch_indices,farthest,:,:]

        centroid_line_dir = centroid_line[:,0,:] - centroid_line[:,1,:]
        centroid_line_norm = centroid_line_dir / torch.norm(centroid_line_dir,dim=-1,keepdim=True)

        dist = (centroid_line_norm[:,None,:] * seg_dir_norm).sum(-1)
        dist = 1 - torch.abs(dist)
        # print(ss)
        # distA= PointToLineSegmentDis(centroid_line[0,0],seg[0,:,0,:],seg[0,:,1,:])
        # distB = PointToLineSegmentDis(centroid_line[0, 1], seg[0, :, 0, :], seg[0, :, 1, :])
        # dist = torch.maximum(distA,distB).unsqueeze(0)


        # centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        # dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    # centroids = torch.cat([centroids*2,centroids*2+1],dim=0)
    # centroids = centroids.view(1,-1)
    return centroids
//  [python]
def farthest_point_sampleV4(xyz, npoint):
    """
    这个应该是基于角度进行FPS
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
        oldidx: [1,M]
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    def dist(center,other):
        pass
    seg = xyz.view(1,-1,2,3)
    device = xyz.device
    B, N, _, C = seg.shape
    if npoint>=N:
        x = np.arange(npoint)%N
        return torch.from_numpy(x).unsqueeze(0)
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device) * 0
    batch_indices = torch.arange(B, dtype=torch.long).to(device)

    seg_dir = seg[:,:,0,:] - seg[:,:,1,:]
    seg_dir_norm = seg_dir / torch.linalg.norm(seg_dir,dim=-1,keepdims=True)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid_line = seg[batch_indices,farthest,:,:]

        centroid_line_dir = centroid_line[:,0,:] - centroid_line[:,1,:]
        centroid_line_norm = centroid_line_dir / torch.linalg.norm(centroid_line_dir,dim=-1,keepdims=True)

        dist = (centroid_line_norm[:,None,:] * seg_dir_norm).sum(-1)
        dist = 1 - torch.abs(dist)
        # print(ss)
        # distA= PointToLineSegmentDis(centroid_line[0,0],seg[0,:,0,:],seg[0,:,1,:])
        # distB = PointToLineSegmentDis(centroid_line[0, 1], seg[0, :, 0, :], seg[0, :, 1, :])
        # dist = torch.maximum(distA,distB).unsqueeze(0)


        # centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        # dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    # centroids = torch.cat([centroids*2,centroids*2+1],dim=0)
    # centroids = centroids.view(1,-1)
    return centroids
//  [python]
def forward(self, x_true, x_pred):
        S_true = self.spec(x_true) + self.eps
        S_pred = self.spec(x_pred) + self.eps
        
        converge_term = torch.mean(torch.norm(S_true - S_pred, dim = (1, 2)) / torch.norm(S_true + S_pred, dim = (1, 2)))
        
        log_term = F.l1_loss(S_true.log(), S_pred.log())

        loss = converge_term + self.alpha * log_term
        return loss
//  [python]
def forward(self, x_true, x_pred):
        S_true = self.spec(x_true) + self.eps
        S_pred = self.spec(x_pred) + self.eps
        
        converge_term = torch.mean(torch.linalg.norm(S_true - S_pred, dim = (1, 2)) / torch.linalg.norm(S_true + S_pred, dim = (1, 2)))
        
        log_term = F.l1_loss(S_true.log(), S_pred.log())

        loss = converge_term + self.alpha * log_term
        return loss
//  [python]
def test_rand_lower():
    it = 10

    for i in range(it):
        N = torch.randint(32, 2048, (1,)).item()
        A_d, A = common.random_sparse_tri(N, N, 0.05, upper=False, include_diag=False)
        A = A + sp.eye(N)*2. # set diagonal to be 2.
        A = A.double()

        A_c = A.to(gpu)

        b = torch.randn(N, dtype=torch.double)
        b_c = b.to(gpu)

        x = A.solve_triangular(False, False, b)
        x_c = A_c.solve_triangular(False, False, b_c)

        print('N', N)
        print('ans diff', torch.norm(x - x_c.cpu()))
        print('cpu', x)
        print('gpu', x_c.cpu())

        assert(torch.allclose(x, x_c.cpu()))
        assert(torch.norm(b-A@x) / torch.norm(b) < 1e-5)
//  [python]
def test_rand_lower():
    it = 10

    for i in range(it):
        N = torch.randint(32, 2048, (1,)).item()
        A_d, A = common.random_sparse_tri(N, N, 0.05, upper=False, include_diag=False)
        A = A + sp.eye(N)*2. # set diagonal to be 2.
        A = A.double()

        A_c = A.to(gpu)

        b = torch.randn(N, dtype=torch.double)
        b_c = b.to(gpu)

        x = A.solve_triangular(False, False, b)
        x_c = A_c.solve_triangular(False, False, b_c)

        print('N', N)
        print('ans diff', tla.norm(x - x_c.cpu()))
        print('cpu', x)
        print('gpu', x_c.cpu())

        assert(torch.allclose(x, x_c.cpu()))
        assert(tla.norm(b-A@x) / tla.norm(b) < 1e-5)
//  [python]
def nme(self,target, preds):
        target = torch.reshape(target, shape=[-1, 98, 2])
        preds = torch.reshape(preds, shape=[-1, 98, 2])

        norm = torch.norm(target[:, 60, :] - target[:, 72, :], dim=-1)

        distance = torch.mean(torch.norm(preds - target, dim=-1), dim=-1) / norm

        nme = torch.mean(distance)

        return nme
//  [python]
def nme(self,target, preds):
        target = torch.reshape(target, shape=[-1, 98, 2])
        preds = torch.reshape(preds, shape=[-1, 98, 2])

        norm = torch.linalg.norm(target[:, 60, :] - target[:, 72, :], dim=-1)

        distance = torch.mean(torch.linalg.norm(preds - target, dim=-1), dim=-1) / norm

        nme = torch.mean(distance)

        return nme
//  [python]
def compute_cos_sim(
        self,
        points: Optional[TensorType["n", 3]],
        feats: TensorType["n", "c"],
        knn_feats: TensorType["n", "k", "c"],
        knn_pts: TensorType["n", "k", 3],
    ) -> TensorType["n", 1]:
        knn_cos_sims = F.cosine_similarity(feats[:, None], knn_feats, dim=-1)  # (n, k)
        if not self.config.distance_weighted_average:
            return knn_cos_sims.mean(-1, keepdim=True)

        knn_pts_dists = torch.norm(points[:, None] - knn_pts, ord=2, dim=-1)
        knn_pts_weights = torch.exp(-knn_pts_dists)  # TODO: tune temperature
        knn_pts_weights = knn_pts_weights / knn_pts_weights.sum(-1, keepdim=True)  # normalized weights
        # use unnormalized weights would down weight bg_cos_sims too much
        knn_weighed_cos_sims = (knn_cos_sims * knn_pts_weights).sum(-1, keepdim=True)
        return knn_weighed_cos_sims
//  [python]
def compute_cos_sim(
        self,
        points: Optional[TensorType["n", 3]],
        feats: TensorType["n", "c"],
        knn_feats: TensorType["n", "k", "c"],
        knn_pts: TensorType["n", "k", 3],
    ) -> TensorType["n", 1]:
        knn_cos_sims = F.cosine_similarity(feats[:, None], knn_feats, dim=-1)  # (n, k)
        if not self.config.distance_weighted_average:
            return knn_cos_sims.mean(-1, keepdim=True)

        knn_pts_dists = torch.linalg.norm(points[:, None] - knn_pts, ord=2, dim=-1)
        knn_pts_weights = torch.exp(-knn_pts_dists)  # TODO: tune temperature
        knn_pts_weights = knn_pts_weights / knn_pts_weights.sum(-1, keepdim=True)  # normalized weights
        # use unnormalized weights would down weight bg_cos_sims too much
        knn_weighed_cos_sims = (knn_cos_sims * knn_pts_weights).sum(-1, keepdim=True)
        return knn_weighed_cos_sims
//  [python]
def test_lower():
    N = 32
    A = sp.eye(N) * 2 - sp.eye(N, k=-1)
    A_c = A.to(gpu)
    it = 10

    for i in range(it):
        x_e = torch.randn(N)
        b = A @ x_e
        b_c = b.to(gpu)
        x_g = A.solve_triangular(False, False, b)
        x_g_c = A_c.solve_triangular(False, False, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(t.norm(b-A@x_g) / t.norm(b) < 1e-5)

    Au = sp.eye(N) - sp.eye(N, k=-1) * 2
    Au_c = Au.to(gpu)
    for i in range(it):
        x_e = torch.randn(N)
        b = Au @ x_e
        b_c = b.to(gpu)

        x_g = Au.solve_triangular(False, True, b)
        x_g_c = Au_c.solve_triangular(False, True, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(t.norm(b-Au@x_g) /t.norm(b) < 1e-5)
//  [python]
def test_lower():
    N = 32
    A = sp.eye(N) * 2 - sp.eye(N, k=-1)
    A_c = A.to(gpu)
    it = 10

    for i in range(it):
        x_e = torch.randn(N)
        b = A @ x_e
        b_c = b.to(gpu)
        x_g = A.solve_triangular(False, False, b)
        x_g_c = A_c.solve_triangular(False, False, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(tla.norm(b-A@x_g) / tla.norm(b) < 1e-5)

    Au = sp.eye(N) - sp.eye(N, k=-1) * 2
    Au_c = Au.to(gpu)
    for i in range(it):
        x_e = torch.randn(N)
        b = Au @ x_e
        b_c = b.to(gpu)

        x_g = Au.solve_triangular(False, True, b)
        x_g_c = Au_c.solve_triangular(False, True, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(tla.norm(b-Au@x_g) /tla.norm(b) < 1e-5)
//  [python]
def test_upper():
    N = 32
    A = sp.eye(N) * 2 - sp.eye(N, k=1)
    A_c = A.to(gpu)
    it = 10

    for i in range(it):
        x_e = torch.randn(N)
        b = A @ x_e
        b_c = b.to(gpu)
        x_g = A.solve_triangular(True, False, b)
        x_g_c = A_c.solve_triangular(True, False, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(t.norm(b-A@x_g) / t.norm(b) < 1e-5)

    Au = sp.eye(N) - sp.eye(N, k=1) * 2
    Au_c = Au.to(gpu)
    for i in range(it):
        x_e = torch.randn(N)
        b = Au @ x_e
        b_c = b.to(gpu)

        x_g = Au.solve_triangular(True, True, b)
        x_g_c = Au_c.solve_triangular(True, True, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(t.norm(b-Au@x_g) /t.norm(b) < 1e-5)
//  [python]
def test_upper():
    N = 32
    A = sp.eye(N) * 2 - sp.eye(N, k=1)
    A_c = A.to(gpu)
    it = 10

    for i in range(it):
        x_e = torch.randn(N)
        b = A @ x_e
        b_c = b.to(gpu)
        x_g = A.solve_triangular(True, False, b)
        x_g_c = A_c.solve_triangular(True, False, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(tla.norm(b-A@x_g) / tla.norm(b) < 1e-5)

    Au = sp.eye(N) - sp.eye(N, k=1) * 2
    Au_c = Au.to(gpu)
    for i in range(it):
        x_e = torch.randn(N)
        b = Au @ x_e
        b_c = b.to(gpu)

        x_g = Au.solve_triangular(True, True, b)
        x_g_c = Au_c.solve_triangular(True, True, b_c)

        assert(torch.allclose(x_g, x_g_c.cpu()))
        assert(tla.norm(b-Au@x_g) /tla.norm(b) < 1e-5)
//  [python]
def outside_sphere_trans(self, pts, pts_norm=None, filtered=False):
        # r^2 = x^2 + y^2 + z^2; x = x / r^2
        out_pts = pts.clone()
        if filtered:
            out_pts = out_pts / pts_norm ** 2
            return out_pts
        if pts_norm is None:
            pts_norm = torch.norm(pts, p=2, dim=-1, keepdim=True)
        inside_sphere = (pts_norm < 1.0)
        out_pts[~inside_sphere[...,0]] = out_pts[~inside_sphere[...,0]] / pts_norm[~inside_sphere[...,0]] ** 2
        out_pts[inside_sphere[...,0]] = -10
        return out_pts, ~inside_sphere
//  [python]
def outside_sphere_trans(self, pts, pts_norm=None, filtered=False):
        # r^2 = x^2 + y^2 + z^2; x = x / r^2
        out_pts = pts.clone()
        if filtered:
            out_pts = out_pts / pts_norm ** 2
            return out_pts
        if pts_norm is None:
            pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True)
        inside_sphere = (pts_norm < 1.0)
        out_pts[~inside_sphere[...,0]] = out_pts[~inside_sphere[...,0]] / pts_norm[~inside_sphere[...,0]] ** 2
        out_pts[inside_sphere[...,0]] = -10
        return out_pts, ~inside_sphere
//  [python]
def batchwise_cosine_similarity(Z,B):
    B = B.T
    Z_norm = torch.norm(Z, dim=1, keepdim=True)  # Size (n, 1).
    B_norm = torch.norm(B, dim=0, keepdim=True)  # Size (1, b).
    cosine_similarity = ((Z @ B) / (Z_norm @ B_norm)).T
    return cosine_similarity
//  [python]
def batchwise_cosine_similarity(Z,B):
    # https://www.h4pz.co/blog/2021/4/2/batch-cosine-similarity-in-pytorch-or-numpy-jax-cupy-etc
    B = B.T
    Z_norm = torch.linalg.norm(Z, dim=1, keepdim=True)  # Size (n, 1).
    B_norm = torch.linalg.norm(B, dim=0, keepdim=True)  # Size (1, b).
    cosine_similarity = ((Z @ B) / (Z_norm @ B_norm)).T
    return cosine_similarity
//  [python]
def batchwise_pearson_correlation(Z, B):
    # Calculate means
    Z_mean = torch.mean(Z, dim=1, keepdim=True)
    B_mean = torch.mean(B, dim=1, keepdim=True)

    # Subtract means
    Z_centered = Z - Z_mean
    B_centered = B - B_mean

    # Calculate Pearson correlation coefficient
    numerator = Z_centered @ B_centered.T
    Z_centered_norm = torch.norm(Z_centered, dim=1, keepdim=True)
    B_centered_norm = torch.norm(B_centered, dim=1, keepdim=True)
    denominator = Z_centered_norm @ B_centered_norm.T

    pearson_correlation = (numerator / denominator)
    return pearson_correlation
//  [python]
def batchwise_pearson_correlation(Z, B):
    # Calculate means
    Z_mean = torch.mean(Z, dim=1, keepdim=True)
    B_mean = torch.mean(B, dim=1, keepdim=True)

    # Subtract means
    Z_centered = Z - Z_mean
    B_centered = B - B_mean

    # Calculate Pearson correlation coefficient
    numerator = Z_centered @ B_centered.T
    Z_centered_norm = torch.linalg.norm(Z_centered, dim=1, keepdim=True)
    B_centered_norm = torch.linalg.norm(B_centered, dim=1, keepdim=True)
    denominator = Z_centered_norm @ B_centered_norm.T

    pearson_correlation = (numerator / denominator)
    return pearson_correlation
//  [python]
def _matrix_unit_norm(A):
    """Scale a matrix to have unit Frobenius norm.

    :param A: a matrix
    :return: the matrix divided by its Frobenius norm
    """
    return A / torch.norm(A)
//  [python]
def _matrix_unit_norm(A):
    """Scale a matrix to have unit Frobenius norm.

    :param A: a matrix
    :return: the matrix divided by its Frobenius norm
    """
    return A / torch.linalg.norm(A, ord='fro')
//  [python]
def _get_spatially_blurred_vert_labels(self, sigma=0.01):
        """
        calculates spatially blurred semantic vertex labels of shape V x C

        uses unmodified template mesh (due to memory constraints) so make sure to upsample result when upsampling mesh

        :return: V x C with scores between 0 and 1.
        """
        UPPER_PLACEHOLDER = 1

        meshes = Meshes(verts=[self._v_template], faces=[self._faces])
        verts = meshes.verts_packed()
        edges = meshes.edges_packed()
        V = verts.shape[0]

        # initializing distance table
        distances = torch.zeros(V, V, dtype=verts.dtype, device=verts.device) + UPPER_PLACEHOLDER
        distances[torch.arange(V), torch.arange(V)] = 0
        e1, e2 = edges.unbind(1)
        distances[e1, e2] = distances[e2, e1] = torch.norm(verts[e1] - verts[e2], dim=-1, p=2)
        max_neighbours = (torch.mean(torch.sum(distances < UPPER_PLACEHOLDER, dim=-1).float() - 1) * 2).int()
        assert torch.max(distances) == UPPER_PLACEHOLDER  # check if upper placeholder was set high enough

        # iteratively propagate distances to linked vertices
        v_helper = torch.arange(V).view(V, 1).expand(V, max_neighbours)
        for i in range(5):
            neighbours = torch.argsort(distances, dim=-1)[:, :max_neighbours]  # V x neighb
            distances = torch.min(distances[neighbours] + distances[v_helper, neighbours].unsqueeze(-1), dim=1).values

        weights = torch.softmax((-(distances ** 2) / (2 * sigma ** 2)), dim=-1)

        vert_labels = torch.zeros(V, len(self._parts), dtype=torch.float, device=self._v_template.device)
        for i, (key, verts) in enumerate(self._parts.items()):
            vert_labels[verts, i] = 1.

        blurred_semantics = torch.matmul(weights, vert_labels)  # (V, C)

        return blurred_semantics
//  [python]
def _get_spatially_blurred_vert_labels(self, sigma=0.01):
        """
        calculates spatially blurred semantic vertex labels of shape V x C

        uses unmodified template mesh (due to memory constraints) so make sure to upsample result when upsampling mesh

        :return: V x C with scores between 0 and 1.
        """
        UPPER_PLACEHOLDER = 1

        meshes = Meshes(verts=[self._v_template], faces=[self._faces])
        verts = meshes.verts_packed()
        edges = meshes.edges_packed()
        V = verts.shape[0]

        # initializing distance table
        distances = torch.zeros(V, V, dtype=verts.dtype, device=verts.device) + UPPER_PLACEHOLDER
        distances[torch.arange(V), torch.arange(V)] = 0
        e1, e2 = edges.unbind(1)
        distances[e1, e2] = distances[e2, e1] = torch.linalg.norm(verts[e1] - verts[e2], dim=-1, ord=2)
        max_neighbours = (torch.mean(torch.sum(distances < UPPER_PLACEHOLDER, dim=-1).float() - 1) * 2).int()
        assert torch.max(distances) == UPPER_PLACEHOLDER  # check if upper placeholder was set high enough

        # iteratively propagate distances to linked vertices
        v_helper = torch.arange(V).view(V, 1).expand(V, max_neighbours)
        for i in range(5):
            neighbours = torch.argsort(distances, dim=-1)[:, :max_neighbours]  # V x neighb
            distances = torch.min(distances[neighbours] + distances[v_helper, neighbours].unsqueeze(-1), dim=1).values

        weights = torch.softmax((-(distances ** 2) / (2 * sigma ** 2)), dim=-1)

        vert_labels = torch.zeros(V, len(self._parts), dtype=torch.float, device=self._v_template.device)
        for i, (key, verts) in enumerate(self._parts.items()):
            vert_labels[verts, i] = 1.

        blurred_semantics = torch.matmul(weights, vert_labels)  # (V, C)

        return blurred_semantics
//  [python]
def _contains_impl(self, pt: Point, atol: float = 1e-6) -> bool:
        # Test shape
        if pt.shape != (self.m, self.m):
            return False
        # Test symmetric
        if not torch.allclose(pt, pt.T, atol=atol):
            return False
        if self._unbiased:
            # Test that the diagonal is zero
            if not torch.allclose(torch.diag(pt), pt.new_zeros((self.m,)), atol=atol):
                return False
        else:
            # Test positive (semi-) definiteness
            e = torch.linalg.eigvalsh(pt)
            if not all(e >= -atol):
                return False
            # Test that the matrix is centered
            if not is_centered(pt, atol=atol):
                return False
        # Test unit Frobenius norm
        if not torch.isclose(torch.norm(pt, ord='fro'), pt.new_ones((1,)), atol=atol):
            return True
//  [python]
def _contains_impl(self, pt: Point, atol: float = 1e-6) -> bool:
        # Test shape
        if pt.shape != (self.m, self.m):
            return False
        # Test symmetric
        if not torch.allclose(pt, pt.T, atol=atol):
            return False
        if self._unbiased:
            # Test that the diagonal is zero
            if not torch.allclose(torch.diag(pt), pt.new_zeros((self.m,)), atol=atol):
                return False
        else:
            # Test positive (semi-) definiteness
            e = torch.linalg.eigvalsh(pt)
            if not all(e >= -atol):
                return False
            # Test that the matrix is centered
            if not is_centered(pt, atol=atol):
                return False
        # Test unit Frobenius norm
        if not torch.isclose(torch.linalg.norm(pt, ord='fro'), pt.new_ones((1,)), atol=atol):
            return False
        return True
//  [python]
def forward(self, x):
    t = self.t.type_as(x)

    if self.training:
      if self.opt['new_edges'] == 'random':
        self.add_random_edges()
      elif self.opt['new_edges'] == 'k_hop':
        self.add_khop_edges(k=2)
      elif self.opt['new_edges'] == 'random_walk' and self.odefunc.attention_weights is not None:
        self.add_rw_edges()



    attention_weights = self.get_attention_weights(x)
    # create attention mask
    if self.training:
      with torch.no_grad():
        mean_att = attention_weights.mean(dim=1, keepdim=False)
        if self.opt['use_flux']:
          src_features = x[self.data_edge_index[0, :], :]
          dst_features = x[self.data_edge_index[1, :], :]
          delta = torch.norm(src_features-dst_features, dim=1)
          mean_att = mean_att * delta
        threshold = torch.quantile(mean_att, 1-self.opt['att_samp_pct'])
        mask = mean_att > threshold
        self.odefunc.edge_index = self.data_edge_index[:, mask.T]
        sampled_attention_weights = self.renormalise_attention(mean_att[mask])
        print('retaining {} of {} edges'.format(self.odefunc.edge_index.shape[1], self.data_edge_index.shape[1]))
        self.odefunc.attention_weights = sampled_attention_weights
    else:
      self.odefunc.edge_index = self.data_edge_index
      self.odefunc.attention_weights = attention_weights.mean(dim=1, keepdim=False)
    self.reg_odefunc.odefunc.edge_index, self.reg_odefunc.odefunc.edge_weight = self.odefunc.edge_index, self.odefunc.edge_weight
    self.reg_odefunc.odefunc.attention_weights = self.odefunc.attention_weights




    integrator = self.train_integrator if self.training else self.test_integrator
    
    reg_states = tuple( torch.zeros(x.size(0)).to(x) for i in range(self.nreg) )

    func = self.reg_odefunc if self.training and self.nreg > 0 else self.odefunc
    state = (x,) + reg_states if self.training and self.nreg > 0 else x

    if self.opt["adjoint"] and self.training:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options=dict(step_size=self.opt['step_size'], max_iters=self.opt['max_iters']),
        adjoint_method=self.opt['adjoint_method'],
        adjoint_options=dict(step_size = self.opt['adjoint_step_size'], max_iters=self.opt['max_iters']),
        atol=self.atol,
        rtol=self.rtol,
        adjoint_atol=self.atol_adjoint,
        adjoint_rtol=self.rtol_adjoint)
    else:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options=dict(step_size=self.opt['step_size'], max_iters=self.opt['max_iters']),
        atol=self.atol,
        rtol=self.rtol)

    if self.training and self.nreg > 0:
      z = state_dt[0][1]
      reg_states = tuple( st[1] for st in state_dt[1:] )
      return z, reg_states
    else: 
      z = state_dt[1]
      return z
//  [python]
def forward(self, x):
    t = self.t.type_as(x)

    if self.training:
      if self.opt['new_edges'] == 'random':
        self.add_random_edges()
      elif self.opt['new_edges'] == 'k_hop':
        self.add_khop_edges(k=2)
      elif self.opt['new_edges'] == 'random_walk' and self.odefunc.attention_weights is not None:
        self.add_rw_edges()



    attention_weights = self.get_attention_weights(x)
    # create attention mask
    if self.training:
      with torch.no_grad():
        mean_att = attention_weights.mean(dim=1, keepdim=False)
        if self.opt['use_flux']:
          src_features = x[self.data_edge_index[0, :], :]
          dst_features = x[self.data_edge_index[1, :], :]
          delta = torch.linalg.norm(src_features-dst_features, dim=1)
          mean_att = mean_att * delta
        threshold = torch.quantile(mean_att, 1-self.opt['att_samp_pct'])
        mask = mean_att > threshold
        self.odefunc.edge_index = self.data_edge_index[:, mask.T]
        sampled_attention_weights = self.renormalise_attention(mean_att[mask])
        print('retaining {} of {} edges'.format(self.odefunc.edge_index.shape[1], self.data_edge_index.shape[1]))
        self.odefunc.attention_weights = sampled_attention_weights
    else:
      self.odefunc.edge_index = self.data_edge_index
      self.odefunc.attention_weights = attention_weights.mean(dim=1, keepdim=False)
    self.reg_odefunc.odefunc.edge_index, self.reg_odefunc.odefunc.edge_weight = self.odefunc.edge_index, self.odefunc.edge_weight
    self.reg_odefunc.odefunc.attention_weights = self.odefunc.attention_weights




    integrator = self.train_integrator if self.training else self.test_integrator
    
    reg_states = tuple( torch.zeros(x.size(0)).to(x) for i in range(self.nreg) )

    func = self.reg_odefunc if self.training and self.nreg > 0 else self.odefunc
    state = (x,) + reg_states if self.training and self.nreg > 0 else x

    if self.opt["adjoint"] and self.training:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options=dict(step_size=self.opt['step_size'], max_iters=self.opt['max_iters']),
        adjoint_method=self.opt['adjoint_method'],
        adjoint_options=dict(step_size = self.opt['adjoint_step_size'], max_iters=self.opt['max_iters']),
        atol=self.atol,
        rtol=self.rtol,
        adjoint_atol=self.atol_adjoint,
        adjoint_rtol=self.rtol_adjoint)
    else:
      state_dt = integrator(
        func, state, t,
        method=self.opt['method'],
        options=dict(step_size=self.opt['step_size'], max_iters=self.opt['max_iters']),
        atol=self.atol,
        rtol=self.rtol)

    if self.training and self.nreg > 0:
      z = state_dt[0][1]
      reg_states = tuple( st[1] for st in state_dt[1:] )
      return z, reg_states
    else: 
      z = state_dt[1]
      return z
//  [python]
def test_clipup_movement(stepsize: float, momentum: float, max_speed: Optional[float]):
    tolerance = 1e-4

    def normalize(step: torch.Tensor) -> torch.Tensor:
        return step / step.norm()

    def clip(update: torch.Tensor, max_speed: float) -> torch.Tensor:
        update_norm = update.norm()
        if update_norm > max_speed:
            update = normalize(update) * max_speed
        return update

    def clipup_step(
        vel: torch.Tensor,
        grad: torch.Tensor,
        stepsize: float,
        momentum: float = 0.9,
        max_speed: Optional[float] = None,
    ) -> torch.Tensor:
        if max_speed is None:
            max_speed = stepsize * 2.0
        raw_update = stepsize * normalize(grad)
        vel = clip(momentum * vel + raw_update, max_speed)
        return vel

    grads = torch.tensor(
        [
            [1, 1, 1],
            [1, 1, 1],
            [1, 1, 0],
            [1, 0, 1],
            [-1, -3, 0],
            [-4, -1, 0],
            [0, 0, -2],
        ],
        dtype=torch.float32,
    )

    num_steps, length = grads.shape
    dtype = grads.dtype

    x0 = torch.zeros(length, dtype=dtype)

    clipup = ClipUp(solution_length=length, dtype=dtype, stepsize=stepsize, momentum=momentum, max_speed=max_speed)

    x = x0.clone()

    desired_vel = torch.zeros(length, dtype=dtype)
    desired_x = x0.clone()

    for t in range(num_steps):
        grad = grads[t, :]

        x = x + clipup.ascent(grad)

        desired_vel = clipup_step(desired_vel, grad, stepsize=stepsize, momentum=momentum, max_speed=max_speed)
        desired_x = desired_x + desired_vel

    assert_allclose(x, desired_x, atol=tolerance)
//  [python]
def test_clipup_movement(stepsize: float, momentum: float, max_speed: Optional[float]):
    tolerance = 1e-4

    def normalize(step: torch.Tensor) -> torch.Tensor:
        return step / torch.linalg.norm(step)

    def clip(update: torch.Tensor, max_speed: float) -> torch.Tensor:
        update_norm = torch.linalg.norm(update)
        if update_norm > max_speed:
            update = normalize(update) * max_speed
        return update

    def clipup_step(
        vel: torch.Tensor,
        grad: torch.Tensor,
        stepsize: float,
        momentum: float = 0.9,
        max_speed: Optional[float] = None,
    ) -> torch.Tensor:
        if max_speed is None:
            max_speed = stepsize * 2.0
        raw_update = stepsize * normalize(grad)
        vel = clip(momentum * vel + raw_update, max_speed)
        return vel

    grads = torch.tensor(
        [
            [1, 1, 1],
            [1, 1, 1],
            [1, 1, 0],
            [1, 0, 1],
            [-1, -3, 0],
            [-4, -1, 0],
            [0, 0, -2],
        ],
        dtype=torch.float32,
    )

    num_steps, length = grads.shape
    dtype = grads.dtype

    x0 = torch.zeros(length, dtype=dtype)

    clipup = ClipUp(solution_length=length, dtype=dtype, stepsize=stepsize, momentum=momentum, max_speed=max_speed)

    x = x0.clone()

    desired_vel = torch.zeros(length, dtype=dtype)
    desired_x = x0.clone()

    for t in range(num_steps):
        grad = grads[t, :]

        x = x + clipup.ascent(grad)

        desired_vel = clipup_step(desired_vel, grad, stepsize=stepsize, momentum=momentum, max_speed=max_speed)
        desired_x = desired_x + desired_vel

    assert_allclose(x, desired_x, atol=tolerance)
//  [python]
def test_subset_active_compute_exponential_decay_function(self):
        a = torch.tensor([1.0, 2.0]).view(2, 1)
        a_p = torch.tensor([3.0, 4.0]).view(2, 1)
        a = torch.cat((a, a_p), 1)
        b = torch.tensor([2.0, 4.0]).view(2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel(active_dims=[0])
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[4.0, 6.0], [5.0, 7.0]])
        actual = offset + diff.pow(-power)
        res = kernel(a, b).to_dense()

        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_subset_active_compute_exponential_decay_function(self):
        a = torch.tensor([1.0, 2.0]).view(2, 1)
        a_p = torch.tensor([3.0, 4.0]).view(2, 1)
        a = torch.cat((a, a_p), 1)
        b = torch.tensor([2.0, 4.0]).view(2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel(active_dims=[0])
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[4.0, 6.0], [5.0, 7.0]])
        actual = offset + diff.pow(-power)
        res = kernel(a, b).to_dense()

        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def test_computes_exponential_decay_function(self):
        a = torch.tensor([1.0, 2.0]).view(2, 1)
        b = torch.tensor([2.0, 4.0]).view(2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel()
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[4.0, 6.0], [5.0, 7.0]])
        actual = offset + torch.tensor([1.0]).div(diff.pow(power))
        res = kernel(a, b).to_dense()

        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_computes_exponential_decay_function(self):
        a = torch.tensor([1.0, 2.0]).view(2, 1)
        b = torch.tensor([2.0, 4.0]).view(2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel()
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[4.0, 6.0], [5.0, 7.0]])
        actual = offset + torch.tensor([1.0]).div(diff.pow(power))
        res = kernel(a, b).to_dense()

        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def exp_map_SE3(tangent_vector: torch.Tensor) -> torch.Tensor:
    # TODO: Modify this to support multiple batch dimensions
    """Compute the exponential map `se(3) -> SE(3)`.

    This can be used for learning pose deltas on `SE(3)`.

    Args:
        tangent_vector: A tangent vector from `se(3)`.

    Returns:
        [R|t] transformation matrices.
    """

    tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)
    tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)

    theta = torch.norm(tangent_vector_ang, dim=1).unsqueeze(1)
    theta2 = theta**2
    theta3 = theta**3

    near_zero = theta < 1e-2
    non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)
    theta_nz = torch.where(near_zero, non_zero, theta)
    theta2_nz = torch.where(near_zero, non_zero, theta2)
    theta3_nz = torch.where(near_zero, non_zero, theta3)

    # Compute the rotation
    sine = theta.sin()
    cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
    sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)
    ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)
    ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)

    ret[:, 0, 0] += cosine.view(-1)
    ret[:, 1, 1] += cosine.view(-1)
    ret[:, 2, 2] += cosine.view(-1)
    temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)
    ret[:, 0, 1] -= temp[:, 2]
    ret[:, 1, 0] += temp[:, 2]
    ret[:, 0, 2] += temp[:, 1]
    ret[:, 2, 0] -= temp[:, 1]
    ret[:, 1, 2] -= temp[:, 0]
    ret[:, 2, 1] += temp[:, 0]

    # Compute the translation
    sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)
    theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)

    ret[:, :, 3:] = sine_by_theta * tangent_vector_lin
    ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)
    ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (
        tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)
    )
    return ret
//  [python]
def exp_map_SE3(tangent_vector: torch.Tensor) -> torch.Tensor:
    # TODO: Modify this to support multiple batch dimensions
    """Compute the exponential map `se(3) -> SE(3)`.

    This can be used for learning pose deltas on `SE(3)`.

    Args:
        tangent_vector: A tangent vector from `se(3)`.

    Returns:
        [R|t] transformation matrices.
    """

    tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)
    tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)

    theta = torch.linalg.norm(tangent_vector_ang, dim=1).unsqueeze(1)
    theta2 = theta**2
    theta3 = theta**3

    near_zero = theta < 1e-2
    non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)
    theta_nz = torch.where(near_zero, non_zero, theta)
    theta2_nz = torch.where(near_zero, non_zero, theta2)
    theta3_nz = torch.where(near_zero, non_zero, theta3)

    # Compute the rotation
    sine = theta.sin()
    cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
    sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)
    ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)
    ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)

    ret[:, 0, 0] += cosine.view(-1)
    ret[:, 1, 1] += cosine.view(-1)
    ret[:, 2, 2] += cosine.view(-1)
    temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)
    ret[:, 0, 1] -= temp[:, 2]
    ret[:, 1, 0] += temp[:, 2]
    ret[:, 0, 2] += temp[:, 1]
    ret[:, 2, 0] -= temp[:, 1]
    ret[:, 1, 2] -= temp[:, 0]
    ret[:, 2, 1] += temp[:, 0]

    # Compute the translation
    sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)
    theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)

    ret[:, :, 3:] = sine_by_theta * tangent_vector_lin
    ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)
    ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (
        tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)
    )
    return ret
//  [python]
def test_initialize_lengthscale(self):
        kernel = ExponentialDecayKernel()
        kernel.initialize(lengthscale=1)
        actual_value = torch.tensor(1.0).view_as(kernel.lengthscale)
        self.assertLess(torch.norm(kernel.lengthscale - actual_value), 1e-5)
//  [python]
def test_initialize_lengthscale(self):
        kernel = ExponentialDecayKernel()
        kernel.initialize(lengthscale=1)
        actual_value = torch.tensor(1.0).view_as(kernel.lengthscale)
        self.assertLess(torch.linalg.norm(kernel.lengthscale - actual_value), 1e-5)
//  [python]
def test_subset_active_exponential_decay_function_batch(self):
        a = torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0], [4.0, 0.0]]).view(2, 2, 2)
        b = torch.tensor([[5.0, 6.0], [7.0, 8.0]]).view(2, 2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]), active_dims=[0])
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        actual = torch.zeros(2, 2, 2)

        diff = torch.tensor([[7.0, 8.0], [8.0, 9.0]])
        actual[0, :, :] = offset + torch.tensor([1.0]).div(diff.pow(power))

        diff = torch.tensor([[11.0, 12.0], [12.0, 13.0]])
        actual[1, :, :] = offset + torch.tensor([1.0]).div(diff.pow(power))

        res = kernel(a, b).to_dense()
        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_subset_active_exponential_decay_function_batch(self):
        a = torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0], [4.0, 0.0]]).view(2, 2, 2)
        b = torch.tensor([[5.0, 6.0], [7.0, 8.0]]).view(2, 2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]), active_dims=[0])
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        actual = torch.zeros(2, 2, 2)

        diff = torch.tensor([[7.0, 8.0], [8.0, 9.0]])
        actual[0, :, :] = offset + torch.tensor([1.0]).div(diff.pow(power))

        diff = torch.tensor([[11.0, 12.0], [12.0, 13.0]])
        actual[1, :, :] = offset + torch.tensor([1.0]).div(diff.pow(power))

        res = kernel(a, b).to_dense()
        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def test_computes_exponential_decay_function_batch(self):
        a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]).view(2, 2, 1)
        b = torch.tensor([[5.0, 6.0], [7.0, 8.0]]).view(2, 2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        actual = torch.zeros(2, 2, 2)

        diff = torch.tensor([[7.0, 8.0], [8.0, 9.0]])
        actual[0, :, :] = offset + diff.pow(-power)

        diff = torch.tensor([[11.0, 12.0], [12.0, 13.0]])
        actual[1, :, :] = offset + diff.pow(-power)

        res = kernel(a, b).to_dense()
        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_computes_exponential_decay_function_batch(self):
        a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]).view(2, 2, 1)
        b = torch.tensor([[5.0, 6.0], [7.0, 8.0]]).view(2, 2, 1)
        lengthscale = 1
        power = 1
        offset = 1

        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        kernel.initialize(lengthscale=lengthscale, power=power, offset=offset)
        kernel.eval()

        actual = torch.zeros(2, 2, 2)

        diff = torch.tensor([[7.0, 8.0], [8.0, 9.0]])
        actual[0, :, :] = offset + diff.pow(-power)

        diff = torch.tensor([[11.0, 12.0], [12.0, 13.0]])
        actual[1, :, :] = offset + diff.pow(-power)

        res = kernel(a, b).to_dense()
        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def test_initialize_lengthscale_batch(self):
        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        ls_init = torch.tensor([1.0, 2.0])
        kernel.initialize(lengthscale=ls_init)
        actual_value = ls_init.view_as(kernel.lengthscale)
        self.assertLess(torch.norm(kernel.lengthscale - actual_value), 1e-5)