jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        # rays_v = torch.matmul(self.poses[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_v = torch.sum(rays_v[..., None, :] * self.poses[img_idx, :3, :3], -1)
        rays_o = self.poses[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda()
//  [python]
def __init__(self, quaternion, translation, rotation, normalize:bool=True, unstack_inputs:bool=False):
    if quaternion is not None:
      assert quaternion.shape[-1] == 4
    if unstack_inputs:
      if rotation is not None:
        rotation = torch.stack([torch.moveaxis(x, -1, 0)   # Unstack.
                    for x in torch.moveaxis(rotation, -2, 0)])  # Unstack.
      translation = torch.moveaxis(translation, -1, 0)  # Unstack.
    if normalize and quaternion is not None:
      quaternion = quaternion / torch.norm(quaternion, dim=-1, keepdim=True)
    if rotation is None:
      rotation = quat_to_rot_pth(quaternion)
    self.quaternion = quaternion
    # self.rotation = [list(row) for row in rotation]
    # self.translation = list(translation)
    self.rotation = rotation
    self.translation = translation
    #assert all(len(row) == 3 for row in self.rotation)
    assert len(self.translation) == 3
//  [python]
def __init__(self, quaternion, translation, rotation, normalize:bool=True, unstack_inputs:bool=False):
    if quaternion is not None:
      assert quaternion.shape[-1] == 4
    if unstack_inputs:
      if rotation is not None:
        rotation = torch.stack([torch.moveaxis(x, -1, 0)   # Unstack.
                    for x in torch.moveaxis(rotation, -2, 0)])  # Unstack.
      translation = torch.moveaxis(translation, -1, 0)  # Unstack.
    if normalize and quaternion is not None:
      quaternion = quaternion / torch.linalg.norm(quaternion, dim=-1, keepdim=True)
    if rotation is None:
      rotation = quat_to_rot_pth(quaternion)
    self.quaternion = quaternion
    # self.rotation = [list(row) for row in rotation]
    # self.translation = list(translation)
    self.rotation = rotation
    self.translation = translation
    #assert all(len(row) == 3 for row in self.rotation)
    assert len(self.translation) == 3
//  [python]
def calculate_ece(p_gt: torch.Tensor, p_pred: torch.Tensor):
    # Calculate Expected Calibration Error (ECE) https://openreview.net/pdf?id=N3FlFslv_J
    # p_gt need to be of shape (-1, 3, #joints)
    # p_pred need to be of shape (-1, n_hypo, 3, #joints)

    # Calculate the median pose per dimension to get a univariate error
    median_pose = torch.quantile(p_pred, q=0.5, dim=1, keepdim=True)

    # Calculate the distance to the mean (skip sorting since we simply count number of elements with a larger error than gt)
    l2_diff = torch.norm(p_pred - median_pose, dim=-2)

    # Calculate the distance between gt and median pose
    l2_diff_gt = torch.norm(p_gt - median_pose[:, 0], dim=-2)

    # Get the empirical quantile by calculating the number of elements above and below the gt distance
    quantile = (l2_diff_gt.unsqueeze(1) < l2_diff).float().mean(dim=1)

    # To follow the paper, the quantiles should be averaged across the GT poses followed by the median along the keypoint dimension
    return quantile
//  [python]
def calculate_ece(p_gt: torch.Tensor, p_pred: torch.Tensor):
    # Calculate Expected Calibration Error (ECE) https://openreview.net/pdf?id=N3FlFslv_J
    # p_gt need to be of shape (-1, 3, #joints)
    # p_pred need to be of shape (-1, n_hypo, 3, #joints)

    # Calculate the median pose per dimension to get a univariate error
    median_pose = torch.quantile(p_pred, q=0.5, dim=1, keepdim=True)

    # Calculate the distance to the mean (skip sorting since we simply count number of elements with a larger error than gt)
    l2_diff = torch.linalg.norm(p_pred - median_pose, dim=-2)

    # Calculate the distance between gt and median pose
    l2_diff_gt = torch.linalg.norm(p_gt - median_pose[:, 0], dim=-2)

    # Get the empirical quantile by calculating the number of elements above and below the gt distance
    quantile = (l2_diff_gt.unsqueeze(1) < l2_diff).float().mean(dim=1)

    # To follow the paper, the quantiles should be averaged across the GT poses followed by the median along the keypoint dimension
    return quantile
//  [python]
def test_no_dim_no_ord_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor().to(device)
        keepdim = random_bool()
        m = torch.norm(input, keepdim=keepdim)
        n = torch.norm(input, keepdim=keepdim)
        return m, n
//  [python]
def test_no_dim_no_ord_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor().to(device)
        keepdim = random_bool()
        m = torch.linalg.norm(input, keepdim=keepdim)
        n = torch.norm(input, keepdim=keepdim)
        return m, n
//  [python]
def test_one_dim_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=4).to(device)
        dim = random(low=0, high=4).to(int)
        k = random().to(float)
        ord = oneof(float("inf"), float("-inf"), k, None)
        keepdim = random_bool()
        m = torch.norm(input, ord, dim, keepdim)
        n = torch.norm(input, ord, dim, keepdim)
        return m, n
//  [python]
def test_one_dim_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=4).to(device)
        dim = random(low=0, high=4).to(int)
        k = random().to(float)
        ord = oneof(float("inf"), float("-inf"), k, None)
        keepdim = random_bool()
        m = torch.linalg.norm(input, ord, dim, keepdim)
        n = torch.norm(input, ord, dim, keepdim)
        return m, n
//  [python]
def test_no_dim_two_shape_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=2).to(device)
        ord = oneof(float("inf"), float("-inf"), "fro", 1, -1)
        keepdim = random().to(bool)
        m = torch.norm(input, ord=ord, keepdim=keepdim)
        return m
//  [python]
def test_no_dim_two_shape_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=2).to(device)
        ord = oneof(float("inf"), float("-inf"), "fro", 1, -1)
        keepdim = random().to(bool)
        m = torch.linalg.norm(input, ord=ord, keepdim=keepdim)
        return m
//  [python]
def test_no_dim_one_shape_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=1).to(device)
        k = random().to(float)
        ord = oneof(float("inf"), float("-inf"), k)
        keepdim = random_bool()
        m = torch.norm(input, p=ord, keepdim=keepdim)
        n = torch.norm(input, p=ord, keepdim=keepdim)
        return m, n
//  [python]
def test_no_dim_one_shape_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=1).to(device)
        k = random().to(float)
        ord = oneof(float("inf"), float("-inf"), k)
        keepdim = random_bool()
        m = torch.linalg.norm(input, ord=ord, keepdim=keepdim)
        n = torch.norm(input, p=ord, keepdim=keepdim)
        return m, n
//  [python]
def test_tuple_dim_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=2).to(device)
        dim = oneof((-2, -1), (0, 1), (-1, 0))
        ord = oneof(float("inf"), float("-inf"), "fro", 1, -1, None)
        keepdim = random().to(bool)
        m = torch.norm(input, ord=ord, dim=dim, keepdim=keepdim)
        return m
//  [python]
def test_tuple_dim_norm_with_random_data(test_case):
        device = random_device()
        input = random_tensor(ndim=2).to(device)
        dim = oneof((-2, -1), (0, 1), (-1, 0))
        ord = oneof(float("inf"), float("-inf"), "fro", 1, -1, None)
        keepdim = random().to(bool)
        m = torch.linalg.norm(input, ord=ord, dim=dim, keepdim=keepdim)
        return m
//  [python]
def forward(self, ligs_coords_pred: List[Tensor], ligs_coords: List[Tensor]) -> Tensor:
        distances = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
                distances.append(torch.norm(lig_coords_pred.mean(dim=0)-lig_coords.mean(dim=0)))
        return torch.tensor(distances).mean()
//  [python]
def forward(self, ligs_coords_pred: List[Tensor], ligs_coords: List[Tensor]) -> Tensor:
        distances = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
                distances.append(torch.linalg.norm(lig_coords_pred.mean(dim=0)-lig_coords.mean(dim=0)))
        return torch.tensor(distances).mean()
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[0, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3

        ###
        ### Code below is borrowed from https://github.com/hjxwhy/mipnerf_pl/blob/master/datasets/datasets.py
        ###
        v = p
        # Distance from each unit-norm direction vector to its x-axis neighbor.
        dx = torch.sqrt(torch.sum((v[:-1, :, :] - v[1:, :, :]) ** 2, -1))
        dx = torch.cat([dx, dx[-2:-1, :]], 0)
        # Cut the distance in half, and then round it out so that it's
        # halfway between inscribed by / circumscribed about the pixel.
        radii = dx[..., None] * 2 / math.sqrt(12) # W, H, 1

        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # W, H, 3
        trans = self.pose_all[idx_0, :3, 3] * (1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose_inv = np.linalg.inv(pose)
        rot = torch.from_numpy(pose_inv[:3, :3]).cuda()
        trans = torch.from_numpy(pose_inv[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3

        # TODO: return intrinsics and pose

        return rays_o.transpose(0, 1), rays_v.transpose(0, 1), radii.transpose(0, 1)
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[0, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3

        ###
        ### Code below is borrowed from https://github.com/hjxwhy/mipnerf_pl/blob/master/datasets/datasets.py
        ###
        v = p
        # Distance from each unit-norm direction vector to its x-axis neighbor.
        dx = torch.sqrt(torch.sum((v[:-1, :, :] - v[1:, :, :]) ** 2, -1))
        dx = torch.cat([dx, dx[-2:-1, :]], 0)
        # Cut the distance in half, and then round it out so that it's
        # halfway between inscribed by / circumscribed about the pixel.
        radii = dx[..., None] * 2 / math.sqrt(12) # W, H, 1

        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        trans = self.pose_all[idx_0, :3, 3] * (1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose_inv = np.linalg.inv(pose)
        rot = torch.from_numpy(pose_inv[:3, :3]).cuda()
        trans = torch.from_numpy(pose_inv[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3

        # TODO: return intrinsics and pose

        return rays_o.transpose(0, 1), rays_v.transpose(0, 1), radii.transpose(0, 1)
//  [python]
def forward(self, ligs_coords_pred: List[Tensor], ligs_coords: List[Tensor]) -> Tensor:
        distances = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
            distances.append(torch.norm(lig_coords_pred.mean(dim=0)-lig_coords.mean(dim=0)))
        return torch.median(torch.tensor(distances))
//  [python]
def forward(self, ligs_coords_pred: List[Tensor], ligs_coords: List[Tensor]) -> Tensor:
        distances = []
        for lig_coords_pred, lig_coords in zip(ligs_coords_pred, ligs_coords):
            distances.append(torch.linalg.norm(lig_coords_pred.mean(dim=0)-lig_coords.mean(dim=0)))
        return torch.median(torch.tensor(distances))
//  [python]
def forward(
        self,
        inputs: Dict[str, torch.Tensor],
    ) -> Dict[str, torch.Tensor]:
        # compute distances of atoms to (0, 0, 0)
        center_dists = torch.norm(inputs[properties.R], dim=1)
        # order by that distance
        _, order = center_dists.sort()
        # store ordered positions and types
        inputs[properties.R] = inputs[properties.R][order]
        inputs[properties.Z] = inputs[properties.Z][order]
        return inputs
//  [python]
def forward(
        self,
        inputs: Dict[str, torch.Tensor],
    ) -> Dict[str, torch.Tensor]:
        # compute distances of atoms to (0, 0, 0)
        center_dists = torch.linalg.norm(inputs[properties.R], dim=1)
        # order by that distance
        _, order = center_dists.sort()
        # store ordered positions and types
        inputs[properties.R] = inputs[properties.R][order]
        inputs[properties.Z] = inputs[properties.Z][order]
        return inputs
//  [python]
def l2_loss(dc1, dc2):
    assert dc1.points is not None
    assert dc2.points is not None
    assert len(dc1.points) == len(dc2.points)
    return torch.norm(dc1.points - dc2.points) / len(dc1.points)
//  [python]
def l2_loss(dc1, dc2):
    assert dc1.points is not None
    assert dc2.points is not None
    assert len(dc1.points) == len(dc2.points)
    return torch.linalg.norm(dc1.points - dc2.points) / len(dc1.points)
//  [python]
a = torch.tensor([1,-2,3],requires_grad=True,dtype=torch.float32)
torch.norm(a,float('inf'))
//  [python]
a = torch.tensor([1,-2,3],requires_grad=True,dtype=torch.float32)
torch.linalg.norm(a,float('inf'))
//  [python]
def gaussian_kernel(X, sigma2=1.0, Y=None, normalized=False, **ignored):
    if normalized:
        X = X / torch.norm(X, dim=1, keepdim=True)
        if Y is not None:
            Y = Y / torch.norm(Y, dim=1, keepdim=True)
    Dxx = compute_pdist_sq(X, Y)
    if sigma2 is None:
        sigma2 = Dxx.median()
    Kx = torch.exp(-Dxx / sigma2)
    return Kx
//  [python]
def gaussian_kernel(X, sigma2=1.0, Y=None, normalized=False, **ignored):
    if normalized:
        X = X / torch.linalg.norm(X, dim=1, keepdim=True)
        if Y is not None:
            Y = Y / torch.linalg.norm(Y, dim=1, keepdim=True)
    Dxx = compute_pdist_sq(X, Y)
    if sigma2 is None:
        sigma2 = Dxx.median()
    Kx = torch.exp(-Dxx / sigma2)
    return Kx
//  [python]
def image_reproj_error(imdata, pts3d, img_ids, entrinsics_dict, intrinsics_dict):
    pts3d_array = torch.ones(max(pts3d.keys()) + 1, 4).cuda()
    for pts_id, pts in tqdm(pts3d.items()):
        pts3d_array[pts_id, :3] = torch.from_numpy(pts.xyz).cuda()

    img_errors = torch.ones(len(img_ids), 1)
    for id, img_id in tqdm(enumerate(img_ids)):
        img = imdata[img_id]
        entrinsic = torch.from_numpy(entrinsics_dict[img_id]).float().cuda()
        intrinsic = torch.from_numpy(intrinsics_dict[img_id]).float().cuda()
        img_p3d = pts3d_array[img.point3D_ids]
        projected = intrinsic @ entrinsic[:3] @ img_p3d.T
        projected[:2, :] /= projected[2, :].unsqueeze(0)
        img_2d = torch.from_numpy(img.xys)
        img_error = torch.norm(projected[:2, :].T.cpu() - img_2d, dim=-1)
        img_errors[id] = torch.sum(img_error) / img_error.size()[0]
        # print(f"img_error: {img_errors[id]}")

    return img_errors
//  [python]
def image_reproj_error(imdata, pts3d, img_ids, entrinsics_dict, intrinsics_dict):
    pts3d_array = torch.ones(max(pts3d.keys()) + 1, 4).cuda()
    for pts_id, pts in tqdm(pts3d.items()):
        pts3d_array[pts_id, :3] = torch.from_numpy(pts.xyz).cuda()

    img_errors = torch.ones(len(img_ids), 1)
    for id, img_id in tqdm(enumerate(img_ids)):
        img = imdata[img_id]
        entrinsic = torch.from_numpy(entrinsics_dict[img_id]).float().cuda()
        intrinsic = torch.from_numpy(intrinsics_dict[img_id]).float().cuda()
        img_p3d = pts3d_array[img.point3D_ids]
        projected = intrinsic @ entrinsic[:3] @ img_p3d.T
        projected[:2, :] /= projected[2, :].unsqueeze(0)
        img_2d = torch.from_numpy(img.xys)
        img_error = torch.linalg.norm(projected[:2, :].T.cpu() - img_2d, dim=-1)
        img_errors[id] = torch.sum(img_error) / img_error.size()[0]
        # print(f"img_error: {img_errors[id]}")

    return img_errors
//  [python]
def test_FrobeniusNorm(self):
        """FloorDiv"""

        class Model(nn.Module):

            def __init__(self):
                super(Model, self).__init__()

            def forward(self, x):
                y = torch.norm(x)
                z = x / y
                return z

        self.trace_and_test([(32,64)], Model(), [self.Desc('float32', -1.0, 1.0)])
//  [python]
def test_FrobeniusNorm(self):
        """FloorDiv"""

        class Model(nn.Module):

            def __init__(self):
                super(Model, self).__init__()

            def forward(self, x):
                y = torch.linalg.norm(x, ord='fro')
                z = x / y
                return z

        self.trace_and_test([(32,64)], Model(), [self.Desc('float32', -1.0, 1.0)])
//  [python]
def test_normalization_layer_forward():
    in_space = R2('x')
    C = Circle(in_space, [1, 0], 2.0)
    model = NormalizationLayer(C)
    ps = RandomUniformSampler(C, n_points=50)
    out = model(ps.sample_points())
    assert isinstance(out, Points)
    assert out.as_tensor.shape == (50, 2)
    assert torch.all(torch.norm(out, dim=1) <=1)
//  [python]
def test_normalization_layer_forward():
    in_space = R2('x')
    C = Circle(in_space, [1, 0], 2.0)
    model = NormalizationLayer(C)
    ps = RandomUniformSampler(C, n_points=50)
    out = model(ps.sample_points())
    assert isinstance(out, Points)
    assert out.as_tensor.shape == (50, 2)
    assert torch.all(torch.linalg.norm(out, dim=1) <=1)
//  [python]
def loss_function(self, recon_x, x, mu, log_var, z):
        BCE =  F.mse_loss(
            recon_x.reshape(x.shape[0], -1), x.reshape(x.shape[0], -1), reduction='none'
        ).sum()
        if self.model_name == 'VAE':
            KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

        elif self.model_name =='AE':
            KLD = torch.zeros(1).to(self.device)

        elif self.model_name =='GeoAE':
            KLD = torch.norm(log_var.exp() - torch.ones_like(log_var), dim=1).sum()

        return BCE + self.beta * KLD, BCE, KLD
//  [python]
def loss_function(self, recon_x, x, mu, log_var, z):
        BCE =  F.mse_loss(
            recon_x.reshape(x.shape[0], -1), x.reshape(x.shape[0], -1), reduction='none'
        ).sum()
        if self.model_name == 'VAE':
            KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

        elif self.model_name =='AE':
            KLD = torch.zeros(1).to(self.device)

        elif self.model_name =='GeoAE':
            KLD = torch.linalg.norm(log_var.exp() - torch.ones_like(log_var), dim=1).sum()

        return BCE + self.beta * KLD, BCE, KLD
//  [python]
def compute_grad_stats(model):
    with torch.no_grad():
        stats = []
        for name, p in get_unwrapped_model_if_wrapped(model).reader.named_parameters():
            if p.grad is not None:
                s1 = torch.min(torch.abs(p.grad)).item()
                s2 = torch.max(torch.abs(p.grad)).item()
                s3 = torch.mean(torch.abs(p.grad)).item()
                s4 = torch.norm(p.grad).item()
                stats += [s1, s2, s3, s4]
            else:
                stats += [0.0, 0.0, 0.0, 0.0]
        stats = torch.Tensor(stats).cuda()
        if torch.distributed.is_initialized():
            torch.distributed.all_reduce(stats)
        stats = stats.view(-1, 4)

        res = {}
        res["skip_example"] = (torch.any(torch.isinf(stats)) or torch.any(torch.isnan(stats))).item()
        res["min"] = stats.min(0)[0][0].item()
        res["max"] = stats.max(0)[0][1].item()
        res["mean"] = stats.mean(0)[2].item()
        return res
//  [python]
def compute_grad_stats(model):
    with torch.no_grad():
        stats = []
        for name, p in get_unwrapped_model_if_wrapped(model).reader.named_parameters():
            if p.grad is not None:
                s1 = torch.min(torch.abs(p.grad)).item()
                s2 = torch.max(torch.abs(p.grad)).item()
                s3 = torch.mean(torch.abs(p.grad)).item()
                s4 = torch.linalg.norm(p.grad).item()
                stats += [s1, s2, s3, s4]
            else:
                stats += [0.0, 0.0, 0.0, 0.0]
        stats = torch.Tensor(stats).cuda()
        if torch.distributed.is_initialized():
            torch.distributed.all_reduce(stats)
        stats = stats.view(-1, 4)

        res = {}
        res["skip_example"] = (torch.any(torch.isinf(stats)) or torch.any(torch.isnan(stats))).item()
        res["min"] = stats.min(0)[0][0].item()
        res["max"] = stats.max(0)[0][1].item()
        res["mean"] = stats.mean(0)[2].item()
        return res
//  [python]
def regularized_min_max_diff(
        output_imgs: Tensor,
        activation_extractor: PytorchActivationExtractor,
        eps: float = 1e-6) -> Tensor:
    """
    Calculate the regularized minimum-maximum difference of output images. We want to maximize
    the difference between the minimum and maximum values of the output, but also to regularize
    their values so that they don't exceed the norm of the last layer's input times the norm of
    the last layer's weights.

    Args:
        output_imgs (Tensor or List[Tensor]): The output of the model on images.
        activation_extractor (PytorchActivationExtractor): The activation extractor for the model.
        eps (float): Small value for numerical stability.

    Returns:
        Tensor: The computed minimum-maximum difference loss.
    """
    # get the input to the last linear layers of the model
    output_layers_inputs = activation_extractor.get_output_layer_input_activation()

    # get the weights of the last linear layers of the model
    weights_output_layers = activation_extractor.get_last_linear_layers_weights()

    if not isinstance(output_imgs, (list, tuple)):
        output_imgs = torch.reshape(output_imgs, [output_imgs.shape[0], output_imgs.shape[1], -1])
        output_imgs = torch.mean(output_imgs, dim=-1)
        output_imgs = [output_imgs]
    output_loss = torch.zeros(1).to(get_working_device())

    for output_weight, output, last_layer_input in zip(weights_output_layers, output_imgs, output_layers_inputs):
        weights_norm = torch.norm(output_weight.squeeze(), dim=1)
        out_max, out_argmax = torch.max(output, dim=1)
        out_min, out_argmin = torch.min(output, dim=1)
        last_layer_avg = torch.reshape(last_layer_input, [last_layer_input.shape[0], last_layer_input.shape[1], -1])
        last_layer_avg = torch.mean(last_layer_avg, dim=-1)
        last_layer_norm = torch.norm(last_layer_avg, dim=1)
        reg_min = torch.abs(torch.abs(out_min) - 0.5 * last_layer_norm * weights_norm[out_argmin])
        reg_max = torch.abs(torch.abs(out_max) - 0.5 * last_layer_norm * weights_norm[out_argmax])
        dynamic_loss = 1 / (out_max - out_min + eps)
        output_loss += torch.mean(reg_min + reg_max + dynamic_loss)
    return output_loss
//  [python]
def regularized_min_max_diff(
        output_imgs: Tensor,
        activation_extractor: PytorchActivationExtractor,
        eps: float = 1e-6) -> Tensor:
    """
    Calculate the regularized minimum-maximum difference of output images. We want to maximize
    the difference between the minimum and maximum values of the output, but also to regularize
    their values so that they don't exceed the norm of the last layer's input times the norm of
    the last layer's weights.

    Args:
        output_imgs (Tensor or List[Tensor]): The output of the model on images.
        activation_extractor (PytorchActivationExtractor): The activation extractor for the model.
        eps (float): Small value for numerical stability.

    Returns:
        Tensor: The computed minimum-maximum difference loss.
    """
    # get the input to the last linear layers of the model
    output_layers_inputs = activation_extractor.get_output_layer_input_activation()

    # get the weights of the last linear layers of the model
    weights_output_layers = activation_extractor.get_last_linear_layers_weights()

    if not isinstance(output_imgs, (list, tuple)):
        output_imgs = torch.reshape(output_imgs, [output_imgs.shape[0], output_imgs.shape[1], -1])
        output_imgs = torch.mean(output_imgs, dim=-1)
        output_imgs = [output_imgs]
    output_loss = torch.zeros(1).to(get_working_device())

    for output_weight, output, last_layer_input in zip(weights_output_layers, output_imgs, output_layers_inputs):
        weights_norm = torch.linalg.norm(output_weight.squeeze(), dim=1)
        out_max, out_argmax = torch.max(output, dim=1)
        out_min, out_argmin = torch.min(output, dim=1)
        last_layer_avg = torch.reshape(last_layer_input, [last_layer_input.shape[0], last_layer_input.shape[1], -1])
        last_layer_avg = torch.mean(last_layer_avg, dim=-1)
        last_layer_norm = torch.linalg.norm(last_layer_avg, dim=1)
        reg_min = torch.abs(torch.abs(out_min) - 0.5 * last_layer_norm * weights_norm[out_argmin])
        reg_max = torch.abs(torch.abs(out_max) - 0.5 * last_layer_norm * weights_norm[out_argmax])
        dynamic_loss = 1 / (out_max - out_min + eps)
        output_loss += torch.mean(reg_min + reg_max + dynamic_loss)
    return output_loss
//  [python]
def contract_to_unisphere(self,
        x: torch.Tensor,
        aabb: torch.Tensor,
        ord: int = 2,
        eps: float = 1e-6,
        derivative: bool = False,
    ):
        aabb_min, aabb_max = torch.split(aabb, 3, dim=-1)
        x = (x - aabb_min) / (aabb_max - aabb_min)
        x = x * 2 - 1  # aabb is at [-1, 1]
        mag = torch.norm(x, ord=ord, dim=-1, keepdim=True)
        mask = mag.squeeze(-1) > 1

        if derivative:
            dev = (2 * mag - 1) / mag**2 + 2 * x**2 * (
                1 / mag**3 - (2 * mag - 1) / mag**4
            )
            dev[~mask] = 1.0
            dev = torch.clamp(dev, min=eps)
            return dev
        else:
            x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask])
            x = x / 4 + 0.5  # [-inf, inf] is at [0, 1]
            return x
//  [python]
def contract_to_unisphere(self,
        x: torch.Tensor,
        aabb: torch.Tensor,
        ord: int = 2,
        eps: float = 1e-6,
        derivative: bool = False,
    ):
        aabb_min, aabb_max = torch.split(aabb, 3, dim=-1)
        x = (x - aabb_min) / (aabb_max - aabb_min)
        x = x * 2 - 1  # aabb is at [-1, 1]
        mag = torch.linalg.norm(x, ord=ord, dim=-1, keepdim=True)
        mask = mag.squeeze(-1) > 1

        if derivative:
            dev = (2 * mag - 1) / mag**2 + 2 * x**2 * (
                1 / mag**3 - (2 * mag - 1) / mag**4
            )
            dev[~mask] = 1.0
            dev = torch.clamp(dev, min=eps)
            return dev
        else:
            x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask])
            x = x / 4 + 0.5  # [-inf, inf] is at [0, 1]
            return x
//  [python]
def contract(
    x: Tensor,
    aabb: Tensor,
    ord: Union[str, int] = None,
) -> Tensor:
    """
    Contract the input tensor to the unit cube using piecewise projective function.
    """
    # similar to the one in MeRF paper
    aabb_min, aabb_max = torch.split(aabb, 3, dim=-1)
    x = (x - aabb_min) / (aabb_max - aabb_min)  # 0~1
    x = x * 2 - 1  # aabb is at [-1, 1]
    mag = torch.norm(x, ord=ord, dim=-1, keepdim=True)
    x = torch.where(mag < 1, x, (2 - 1 / mag) * (x / mag))
    x = x / 4 + 0.5  # [-inf, inf] is at [0, 1]
    return x
//  [python]
def contract(
    x: Tensor,
    aabb: Tensor,
    ord: Union[str, int] = None,
) -> Tensor:
    """
    Contract the input tensor to the unit cube using piecewise projective function.
    """
    # similar to the one in MeRF paper
    aabb_min, aabb_max = torch.split(aabb, 3, dim=-1)
    x = (x - aabb_min) / (aabb_max - aabb_min)  # 0~1
    x = x * 2 - 1  # aabb is at [-1, 1]
    mag = torch.linalg.norm(x, ord=ord, dim=-1, keepdim=True)
    x = torch.where(mag < 1, x, (2 - 1 / mag) * (x / mag))
    x = x / 4 + 0.5  # [-inf, inf] is at [0, 1]
    return x
//  [python]
def recompute_correspondence_scores(self, ref_corr_points, src_corr_points, corr_scores, estimated_transform):
        aligned_src_corr_points = apply_transform(src_corr_points, estimated_transform)
        corr_residuals = torch.norm(ref_corr_points - aligned_src_corr_points, dim=1)
        inlier_masks = torch.lt(corr_residuals, self.acceptance_radius)
        new_corr_scores = corr_scores * inlier_masks.float()
        return new_corr_scores
//  [python]
def recompute_correspondence_scores(self, ref_corr_points, src_corr_points, corr_scores, estimated_transform):
        aligned_src_corr_points = apply_transform(src_corr_points, estimated_transform)
        corr_residuals = torch.linalg.norm(ref_corr_points - aligned_src_corr_points, dim=1)
        inlier_masks = torch.lt(corr_residuals, self.acceptance_radius)
        new_corr_scores = corr_scores * inlier_masks.float()
        return new_corr_scores
//  [python]
def __init__(self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if((rot_mats is None and quats is None) or 
            (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

        if((rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or 
            (quats is not None and quats.shape[-1] != 4)):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if False:
            if(quats is not None):
                quats = quats.type(torch.float32)
            if(rot_mats is not None):
                rot_mats = rot_mats.type(torch.float32)

        if(quats is not None and normalize_quats):
            quats = quats / torch.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def __init__(self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if((rot_mats is None and quats is None) or 
            (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

        if((rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or 
            (quats is not None and quats.shape[-1] != 4)):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if False:
            if(quats is not None):
                quats = quats.type(torch.float32)
            if(rot_mats is not None):
                rot_mats = rot_mats.type(torch.float32)

        if(quats is not None and normalize_quats):
            quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Converts the XY base velocity command to arrow direction rotation."""
        # obtain default scale of the marker
        default_scale = self.base_vel_goal_visualizer.cfg.markers["arrow"].scale
        # arrow-scale
        arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
        arrow_scale[:, 0] *= torch.norm(xy_velocity, dim=1) * 3.0
        # arrow-direction
        heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
        zeros = torch.zeros_like(heading_angle)
        arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle)
        # convert everything back from base to world frame
        base_quat_w = self.robot.data.root_quat_w
        arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)

        return arrow_scale, arrow_quat
//  [python]
def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Converts the XY base velocity command to arrow direction rotation."""
        # obtain default scale of the marker
        default_scale = self.base_vel_goal_visualizer.cfg.markers["arrow"].scale
        # arrow-scale
        arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
        arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0
        # arrow-direction
        heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
        zeros = torch.zeros_like(heading_angle)
        arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle)
        # convert everything back from base to world frame
        base_quat_w = self.robot.data.root_quat_w
        arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)

        return arrow_scale, arrow_quat
//  [python]
def calc_distogram(pos, min_bin, max_bin, num_bins):
    dists_2d = torch.norm(
        pos[:, :, None, :] - pos[:, None, :, :], dim=-1)[..., None]
    lower = torch.linspace(
        min_bin,
        max_bin,
        num_bins,
        device=pos.device)
    upper = torch.cat([lower[1:], lower.new_tensor([1e8])], dim=-1)
    dgram = ((dists_2d > lower) * (dists_2d < upper)).type(pos.dtype)
    return dgram
//  [python]
def calc_distogram(pos, min_bin, max_bin, num_bins):
    dists_2d = torch.linalg.norm(
        pos[:, :, None, :] - pos[:, None, :, :], axis=-1)[..., None]
    lower = torch.linspace(
        min_bin,
        max_bin,
        num_bins,
        device=pos.device)
    upper = torch.cat([lower[1:], lower.new_tensor([1e8])], dim=-1)
    dgram = ((dists_2d > lower) * (dists_2d < upper)).type(pos.dtype)
    return dgram
//  [python]
def rotation_matrix(a: Float[Tensor, "3"], b: Float[Tensor, "3"]) -> Float[Tensor, "3 3"]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / t.norm(a)
    b = b / t.norm(b)
    v = torch.cross(a, b)
    c = torch.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (torch.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = t.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
//  [python]
def rotation_matrix(a: Float[Tensor, "3"], b: Float[Tensor, "3"]) -> Float[Tensor, "3 3"]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / torch.linalg.norm(a)
    b = b / torch.linalg.norm(b)
    v = torch.cross(a, b)
    c = torch.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (torch.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = torch.linalg.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
//  [python]
def norm(x: ArrayLike, ord=None, axis=None, keepdims: KeepDims = False):
    x = _atleast_float_1(x)
    return torch.norm(x, ord=ord, dim=axis)
//  [python]
def norm(x: ArrayLike, ord=None, axis=None, keepdims: KeepDims = False):
    x = _atleast_float_1(x)
    return torch.linalg.norm(x, ord=ord, dim=axis)
//  [python]
def quat_to_rotvec(quat, eps=1e-6):
    # w > 0 to ensure 0 <= angle <= pi
    flip = (quat[..., :1] < 0).float()
    quat = (-1 * quat) * flip + (1 - flip) * quat

    angle = 2 * torch.atan2(
        torch.norm(quat[..., 1:], dim=-1),
        quat[..., 0]
    )

    angle2 = angle * angle
    small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
    large_angle_scales = angle / torch.sin(angle / 2 + eps)

    small_angles = (angle <= 1e-3).float()
    rot_vec_scale = small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
    rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
    return rot_vec
//  [python]
def quat_to_rotvec(quat, eps=1e-6):
    # w > 0 to ensure 0 <= angle <= pi
    flip = (quat[..., :1] < 0).float()
    quat = (-1 * quat) * flip + (1 - flip) * quat

    angle = 2 * torch.atan2(
        torch.linalg.norm(quat[..., 1:], dim=-1),
        quat[..., 0]
    )

    angle2 = angle * angle
    small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
    large_angle_scales = angle / torch.sin(angle / 2 + eps)

    small_angles = (angle <= 1e-3).float()
    rot_vec_scale = small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
    rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
    return rot_vec
//  [python]
def embed_utterance_batch(wavs, using_partials=True, return_partials=False, **kwargs):
    # This torch version is designed to cope with a batch of same lengths wavs
    if not using_partials:
        print(wavs.shape)
        frames = audio.wav_to_mel_spectrogram_batch(wavs)
        embeds = embed_frames_batch(frames)
        if return_partials:
            return embeds, None, None
        return embeds

    wave_slices, mel_slices = compute_partial_slices(wavs.shape[-1], **kwargs)
    max_wave_length = wave_slices[-1].stop
    if max_wave_length >= wavs.shape[-1]:
        wavs = torch.cat([wavs, torch.ones((wavs.shape[0], max_wave_length - wavs.shape[-1]), 
                                            dtype=wavs.dtype, device=wavs.device)], 1)

    frames = audio.wav_to_mel_spectrogram_batch(wavs)
    frames_batch = []
    for i in range(len(frames)):
        frames_batch += [frames[i][s] for s in mel_slices]
    frames_batch = torch.stack(frames_batch, 0)
    partial_embeds = embed_frames_batch(frames_batch, use_torch=True)
    partial_embeds = partial_embeds.view(wavs.shape[0], len(mel_slices), -1)

    raw_embeds = torch.mean(partial_embeds, axis=1, keepdims=False)
    embeds = raw_embeds / torch.norm(raw_embeds, dim=-1, keepdim=True)

    if return_partials:
        return embeds, partial_embeds, wave_slices
    return embeds
//  [python]
def embed_utterance_batch(wavs, using_partials=True, return_partials=False, **kwargs):
    # This torch version is designed to cope with a batch of same lengths wavs
    if not using_partials:
        print(wavs.shape)
        frames = audio.wav_to_mel_spectrogram_batch(wavs)
        embeds = embed_frames_batch(frames)
        if return_partials:
            return embeds, None, None
        return embeds

    wave_slices, mel_slices = compute_partial_slices(wavs.shape[-1], **kwargs)
    max_wave_length = wave_slices[-1].stop
    if max_wave_length >= wavs.shape[-1]:
        wavs = torch.cat([wavs, torch.ones((wavs.shape[0], max_wave_length - wavs.shape[-1]), 
                                            dtype=wavs.dtype, device=wavs.device)], 1)

    frames = audio.wav_to_mel_spectrogram_batch(wavs)
    frames_batch = []
    for i in range(len(frames)):
        frames_batch += [frames[i][s] for s in mel_slices]
    frames_batch = torch.stack(frames_batch, 0)
    partial_embeds = embed_frames_batch(frames_batch, use_torch=True)
    partial_embeds = partial_embeds.view(wavs.shape[0], len(mel_slices), -1)

    raw_embeds = torch.mean(partial_embeds, axis=1, keepdims=False)
    embeds = raw_embeds / torch.linalg.norm(raw_embeds, axis=-1, keepdims=True)

    if return_partials:
        return embeds, partial_embeds, wave_slices
    return embeds
//  [python]
def forward(self, q_xyzs, k_xyzs, knn_idx, mask):
        _, _, k = knn_idx.shape
        knn_xyzs = index_points(k_xyzs, knn_idx)
        repeated_xyzs = q_xyzs[..., None].repeat(1, 1, 1, k)

        offset_xyzs = knn_xyzs - repeated_xyzs
        direction = F.normalize(offset_xyzs, p=2, dim=1)
        distance = torch.norm(offset_xyzs, dim=1, keepdim=True)
        local_pattern = torch.cat((direction, distance), dim=1)

        position_embedding_expanded = self.embed_positions(local_pattern)

        attn = self.attention(position_embedding_expanded)
        mask_value = -(torch.finfo(attn.dtype).max)
        attn.masked_fill_(~mask[:, None], mask_value)
        attn = F.softmax(attn, dim=-1)
        position_embedding = (position_embedding_expanded * attn).sum(dim=-1)

        return position_embedding
//  [python]
def forward(self, q_xyzs, k_xyzs, knn_idx, mask):
        # q_xyzs: (b, 3, m)
        # k_xyzs: (b, 3, n)
        # knn_idx: (b, m, k)
        # mask: (b, m, k)
        # knn_xyzs: (b, 3, m, k)
        # repeated_xyzs: (b, 3, m, k)
        # direction: (b, 3, m, k)
        # distance: (b, 1, m, k)
        # local_pattern: (b, 4, m, k)
        # position_embedding_expanded: (b, c, m, k)
        # attn: (b, c, m, k)
        # position_embedding: (b, c, m)

        # "query" (q_xyzs) points are the centroids of each point group.
        # "key" (k_xyzs) points are points in the neighborhood of each centroid.

        _, _, k = knn_idx.shape
        knn_xyzs = index_points(k_xyzs, knn_idx)
        repeated_xyzs = q_xyzs[..., None].repeat(1, 1, 1, k)

        # Represent points within a point group as (direction, distance)
        # of offsets from the group centroid.
        offset_xyzs = knn_xyzs - repeated_xyzs
        direction = F.normalize(offset_xyzs, p=2, dim=1)
        distance = torch.linalg.norm(offset_xyzs, dim=1, keepdim=True)
        local_pattern = torch.cat((direction, distance), dim=1)

        # Apply a pointwise MLP to each point.
        position_embedding_expanded = self.embed_positions(local_pattern)

        # Compute self-attention, ignoring points that are not in the
        # neighborhood of the centroid.
        attn = self.attention(position_embedding_expanded)
        mask_value = -(torch.finfo(attn.dtype).max)
        attn.masked_fill_(~mask[:, None], mask_value)
        attn = F.softmax(attn, dim=-1)
        position_embedding = (position_embedding_expanded * attn).sum(dim=-1)

        return position_embedding
//  [python]
def gen_rays(self):
        tx = torch.linspace(0, self.W - 1, self.W)
        ty = torch.linspace(0, self.H - 1, self.H)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[:, None, None, :3, :3], p[None, :, :, :, None]).squeeze() # n, W, H, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # n, W, H, 3
        rays_v = torch.matmul(self.pose_all[:, None, None, :3, :3],  rays_v[:, :, :, :, None]).squeeze()  # n, W, H, 3
        rays_o = self.pose_all[:, None, None, :3, 3].expand(rays_v.shape)  # n, W, H, 3
        return rays_o.transpose(1, 2), rays_v.transpose(1, 2)
//  [python]
def gen_rays(self):
        tx = torch.linspace(0, self.W - 1, self.W)
        ty = torch.linspace(0, self.H - 1, self.H)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[:, None, None, :3, :3], p[None, :, :, :, None]).squeeze() # n, W, H, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # n, W, H, 3
        rays_v = torch.matmul(self.pose_all[:, None, None, :3, :3],  rays_v[:, :, :, :, None]).squeeze()  # n, W, H, 3
        rays_o = self.pose_all[:, None, None, :3, 3].expand(rays_v.shape)  # n, W, H, 3
        return rays_o.transpose(1, 2), rays_v.transpose(1, 2)
//  [python]
def gen_rays_at(self, img_idx, pose = None, resolution_level=1):
        pose_cur = self.get_pose(img_idx, pose)

        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.matmul(pose_cur[None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = pose_cur[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)