jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def disturb_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
        bsz, dim = profile.shape[0], profile.shape[-1]
        profile_norm = torch.norm(profile, dim=-1, keepdim=False)
        spk_count = binary_labels.sum(dim=1)
        prob = np.random.rand(bsz)
        batch_indices = np.nonzero(prob < self.disturb_aug_prob)[0]
        for idx in batch_indices:
            pos_spk_idx = torch.nonzero(spk_count[idx] * mask[idx])
            valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx])
            if len(pos_spk_idx) == 0 or len(valid_spk_idx) == 0:
                continue
            to_disturb_idx = pos_spk_idx[torch.randint(len(pos_spk_idx), ())]
            disturb_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())]
            alpha = self.disturb_alpha * torch.rand(()).item()
            profile[idx, to_disturb_idx] = ((1 - alpha) * profile[idx, to_disturb_idx]
                                            + alpha * profile[idx, disturb_idx])
            profile[idx, to_disturb_idx] = F.normalize(profile[idx, to_disturb_idx], dim=-1)
            mask[idx, to_disturb_idx] = 0

        return profile, binary_labels, mask
//  [python]
def disturb_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
        bsz, dim = profile.shape[0], profile.shape[-1]
        profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False)
        spk_count = binary_labels.sum(dim=1)
        prob = np.random.rand(bsz)
        batch_indices = np.nonzero(prob < self.disturb_aug_prob)[0]
        for idx in batch_indices:
            pos_spk_idx = torch.nonzero(spk_count[idx] * mask[idx])
            valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx])
            if len(pos_spk_idx) == 0 or len(valid_spk_idx) == 0:
                continue
            to_disturb_idx = pos_spk_idx[torch.randint(len(pos_spk_idx), ())]
            disturb_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())]
            alpha = self.disturb_alpha * torch.rand(()).item()
            profile[idx, to_disturb_idx] = ((1 - alpha) * profile[idx, to_disturb_idx]
                                            + alpha * profile[idx, disturb_idx])
            profile[idx, to_disturb_idx] = F.normalize(profile[idx, to_disturb_idx], dim=-1)
            mask[idx, to_disturb_idx] = 0

        return profile, binary_labels, mask
//  [python]
def update(self, joints_rst: Tensor, joints_ref: Tensor,
               lengths: List[int]):
        
        assert joints_rst.shape == joints_ref.shape
        assert joints_rst.dim() == 4
        # (bs, seq, njoint=22, 3)

        self.count += sum(lengths)
        self.count_seq += len(lengths)

        rst = torch.flatten(joints_rst, start_dim=2)
        ref = torch.flatten(joints_ref, start_dim=2)
        
        for i, l in enumerate(lengths):
            if self.task == "pred":
                pred_start = int(l*self.cfg.ABLATION.predict_ratio)
                diff = rst[i,pred_start:] - ref[i,pred_start:]
            elif self.task == "inbetween":
                inbetween_start = int(l*self.cfg.ABLATION.inbetween_ratio)
                inbetween_end = l - int(l*self.cfg.ABLATION.inbetween_ratio)
                diff = rst[i,inbetween_start:inbetween_end] - ref[i,inbetween_start:inbetween_end]
            else:
                print(f"Task {self.task} not implemented.")
                diff = rst - ref
            
            dist = torch.norm(diff, dim=-1)[None]

            ade = dist.mean(dim=1)
            fde = dist[:,-1]
            self.ADE = self.ADE + ade
            self.FDE = self.FDE + fde
//  [python]
def update(self, joints_rst: Tensor, joints_ref: Tensor,
               lengths: List[int]):
        
        assert joints_rst.shape == joints_ref.shape
        assert joints_rst.dim() == 4
        # (bs, seq, njoint=22, 3)

        self.count += sum(lengths)
        self.count_seq += len(lengths)

        rst = torch.flatten(joints_rst, start_dim=2)
        ref = torch.flatten(joints_ref, start_dim=2)
        
        for i, l in enumerate(lengths):
            if self.task == "pred":
                pred_start = int(l*self.cfg.ABLATION.predict_ratio)
                diff = rst[i,pred_start:] - ref[i,pred_start:]
            elif self.task == "inbetween":
                inbetween_start = int(l*self.cfg.ABLATION.inbetween_ratio)
                inbetween_end = l - int(l*self.cfg.ABLATION.inbetween_ratio)
                diff = rst[i,inbetween_start:inbetween_end] - ref[i,inbetween_start:inbetween_end]
            else:
                print(f"Task {self.task} not implemented.")
                diff = rst - ref
            
            dist = torch.linalg.norm(diff, dim=-1)[None]

            ade = dist.mean(dim=1)
            fde = dist[:,-1]
            self.ADE = self.ADE + ade
            self.FDE = self.FDE + fde
//  [python]
def _compute_fine_loss_l2(self, expec_f, expec_f_gt):
        """
        Args:
            expec_f (torch.Tensor): [M, 2] <x, y>
            expec_f_gt (torch.Tensor): [M, 2] <x, y>
        """
        correct_mask = torch.norm(expec_f_gt, p=float('inf'), dim=1) < self.correct_thr
        if correct_mask.sum() == 0:
            if self.training:  # this seldomly happen when training, since we pad prediction with gt
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
            else:
                return None
        offset_l2 = ((expec_f_gt[correct_mask] - expec_f[correct_mask]) ** 2).sum(-1)
        return offset_l2.mean()
//  [python]
def _compute_fine_loss_l2(self, expec_f, expec_f_gt):
        """
        Args:
            expec_f (torch.Tensor): [M, 2] <x, y>
            expec_f_gt (torch.Tensor): [M, 2] <x, y>
        """
        correct_mask = torch.linalg.norm(expec_f_gt, ord=float('inf'), dim=1) < self.correct_thr
        if correct_mask.sum() == 0:
            if self.training:  # this seldomly happen when training, since we pad prediction with gt
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
            else:
                return None
        offset_l2 = ((expec_f_gt[correct_mask] - expec_f[correct_mask]) ** 2).sum(-1)
        return offset_l2.mean()
//  [python]
def check(self, value):
        tol = torch.finfo(value.dtype).eps * value.size(-1) * 10  # 10 is an adjustable fudge factor
        row_norm = torch.norm(value.detach(), dim=-1)
        unit_row_norm = (row_norm - 1.).abs().le(tol).all(dim=-1)
        return _LowerCholesky().check(value) & unit_row_norm
//  [python]
def check(self, value):
        tol = torch.finfo(value.dtype).eps * value.size(-1) * 10  # 10 is an adjustable fudge factor
        row_norm = torch.linalg.norm(value.detach(), dim=-1)
        unit_row_norm = (row_norm - 1.).abs().le(tol).all(dim=-1)
        return _LowerCholesky().check(value) & unit_row_norm
//  [python]
def _compute_fine_loss_l2_std(self, expec_f, expec_f_gt):
        """
        Args:
            expec_f (torch.Tensor): [M, 3] <x, y, std>
            expec_f_gt (torch.Tensor): [M, 2] <x, y>
        """
        # correct_mask tells you which pair to compute fine-loss
        correct_mask = torch.norm(expec_f_gt, float('inf'), dim=1) < self.correct_thr

        # use std as weight that measures uncertainty
        std = expec_f[:, 2]
        inverse_std = 1. / torch.clamp(std, min=1e-10)
        weight = (inverse_std / torch.mean(inverse_std)).detach()  # avoid minizing loss through increase std

        # corner case: no correct coarse match found
        if not correct_mask.any():
            if self.training:  # this seldomly happen during training, since we pad prediction with gt
                               # sometimes there is not coarse-level gt at all.
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
                weight[0] = 0.
            else:
                return None

        # l2 loss with std
        offset_l2 = ((expec_f_gt[correct_mask] - expec_f[correct_mask, :2]) ** 2).sum(-1)
        loss = (offset_l2 * weight[correct_mask]).mean()

        return loss
//  [python]
def _compute_fine_loss_l2_std(self, expec_f, expec_f_gt):
        """
        Args:
            expec_f (torch.Tensor): [M, 3] <x, y, std>
            expec_f_gt (torch.Tensor): [M, 2] <x, y>
        """
        # correct_mask tells you which pair to compute fine-loss
        correct_mask = torch.linalg.norm(expec_f_gt, ord=float('inf'), dim=1) < self.correct_thr

        # use std as weight that measures uncertainty
        std = expec_f[:, 2]
        inverse_std = 1. / torch.clamp(std, min=1e-10)
        weight = (inverse_std / torch.mean(inverse_std)).detach()  # avoid minizing loss through increase std

        # corner case: no correct coarse match found
        if not correct_mask.any():
            if self.training:  # this seldomly happen during training, since we pad prediction with gt
                               # sometimes there is not coarse-level gt at all.
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
                weight[0] = 0.
            else:
                return None

        # l2 loss with std
        offset_l2 = ((expec_f_gt[correct_mask] - expec_f[correct_mask, :2]) ** 2).sum(-1)
        loss = (offset_l2 * weight[correct_mask]).mean()

        return loss
//  [python]
def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor],
        query_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.FloatTensor]:
        image_class_embeds = self.dense0(image_embeds)
        if query_embeds is None:
            device = image_class_embeds.device
            batch_size, num_patches = image_class_embeds.shape[:2]
            pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)
            return (pred_logits, image_class_embeds)

        # Normalize image and text features
        image_class_embeds = image_class_embeds / (torch.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6)
        query_embeds = query_embeds / (torch.norm(query_embeds, dim=-1, keepdim=True) + 1e-6)

        # Get class predictions
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        # Apply a learnable shift and scale to logits
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)
//  [python]
def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor],
        query_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.FloatTensor]:
        image_class_embeds = self.dense0(image_embeds)
        if query_embeds is None:
            device = image_class_embeds.device
            batch_size, num_patches = image_class_embeds.shape[:2]
            pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)
            return (pred_logits, image_class_embeds)

        # Normalize image and text features
        image_class_embeds = image_class_embeds / (torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6)
        query_embeds = query_embeds / (torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6)

        # Get class predictions
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        # Apply a learnable shift and scale to logits
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)
//  [python]
def cdist(X, Y, p=2.0):
    # 2d P, C = X.shape| R, C = Y.shape -> P,R
    P, C = X.shape
    R, C = Y.shape
    # 3d B, P, C = X.shape|1, R, C = Y.shape -> B, P,R
    # D = paddle.linalg.norm(X[:, None, :]-Y[None, :, :], axis=-1)
    """D = paddle.zeros((P, R))
    for i in range(P):
        D[i,:] = paddle.linalg.norm(X[i, None, :]-Y, axis=-1)
        #D[i,:] = (X[i, None, :]-Y).square().sum(-1).sqrt_()
    #"""
    D = []
    for i in range(P):
        D.append(torch.norm(X[i, None, :] - Y, axis=-1))
    D = torch.stack(D, 0)
    return D
//  [python]
def cdist(X, Y, p=2.0):
    # 2d P, C = X.shape| R, C = Y.shape -> P,R
    P, C = X.shape
    R, C = Y.shape
    # 3d B, P, C = X.shape|1, R, C = Y.shape -> B, P,R
    # D = paddle.linalg.norm(X[:, None, :]-Y[None, :, :], axis=-1)
    """D = paddle.zeros((P, R))
    for i in range(P):
        D[i,:] = paddle.linalg.norm(X[i, None, :]-Y, axis=-1)
        #D[i,:] = (X[i, None, :]-Y).square().sum(-1).sqrt_()
    #"""
    D = []
    for i in range(P):
        D.append(torch.linalg.norm(X[i, None, :] - Y, axis=-1))
    D = torch.stack(D, 0)
    return D
//  [python]
def _compute_norm_for_weight(weight, scheme, ord):
        """Compute and return operator norm for given matrix."""
        # fold into matrix operator
        weight = scheme.fold(weight.detach())

        # get and return operator norm of flattened weight
        return torch.norm(weight, p=ord)
//  [python]
def _compute_norm_for_weight(weight, scheme, ord):
        """Compute and return operator norm for given matrix."""
        # fold into matrix operator
        weight = scheme.fold(weight.detach())

        # get and return operator norm of flattened weight
        return torch.linalg.norm(weight, ord=ord)
//  [python]
def extract_udf_mesh(self, world_space=False, resolution=256, dist_threshold_ratio=1.0):
        if self.model_type == 'udf':
            func = self.udf_network_fine.udf

            def func_grad(xyz):
                gradients = self.udf_network_fine.gradient(xyz)
                gradients_mag = torch.norm(gradients, p=2, dim=-1, keepdim=True)
                gradients_norm = gradients / (gradients_mag + 1e-5)  # normalize to unit vector
                return gradients_norm

        elif self.model_type == 'neus':
            func = lambda pts: torch.abs(self.sdf_network_fine.sdf(pts))
            func_grad = self.sdf_network_fine.gradient

        try:
            pred_v, pred_f, pred_mesh, samples, indices = get_mesh_udf_fast(func, func_grad, samples=None,
                                                                            indices=None, N_MC=resolution,
                                                                            gradient=True, eps=0.005,
                                                                            border_gradients=True,
                                                                            smooth_borders=True,
                                                                            dist_threshold_ratio=dist_threshold_ratio)
        except:
            pred_v, pred_f, pred_mesh, samples, indices = get_mesh_udf_fast(func, func_grad, samples=None,
                                                                            indices=None, N_MC=resolution,
                                                                            gradient=True, eps=0.005,
                                                                            border_gradients=False,
                                                                            smooth_borders=False,
                                                                            dist_threshold_ratio=dist_threshold_ratio)

        vertices, triangles = pred_mesh.vertices, pred_mesh.faces
        if world_space:
            vertices = vertices * self.dataset.scale_mats_np[0][0, 0] + self.dataset.scale_mats_np[0][:3, 3][None]

        mesh = trimesh.Trimesh(vertices, triangles)

        os.makedirs(os.path.join(self.base_exp_dir, 'udf_meshes'), exist_ok=True)
        mesh.export(
            os.path.join(self.base_exp_dir, 'udf_meshes', 'udf_res{}_step{}.ply'.format(resolution, self.iter_step)))
//  [python]
def extract_udf_mesh(self, world_space=False, resolution=256, dist_threshold_ratio=1.0):
        if self.model_type == 'udf':
            func = self.udf_network_fine.udf

            def func_grad(xyz):
                gradients = self.udf_network_fine.gradient(xyz)
                gradients_mag = torch.linalg.norm(gradients, ord=2, dim=-1, keepdim=True)
                gradients_norm = gradients / (gradients_mag + 1e-5)  # normalize to unit vector
                return gradients_norm

        elif self.model_type == 'neus':
            func = lambda pts: torch.abs(self.sdf_network_fine.sdf(pts))
            func_grad = self.sdf_network_fine.gradient

        try:
            pred_v, pred_f, pred_mesh, samples, indices = get_mesh_udf_fast(func, func_grad, samples=None,
                                                                            indices=None, N_MC=resolution,
                                                                            gradient=True, eps=0.005,
                                                                            border_gradients=True,
                                                                            smooth_borders=True,
                                                                            dist_threshold_ratio=dist_threshold_ratio)
        except:
            pred_v, pred_f, pred_mesh, samples, indices = get_mesh_udf_fast(func, func_grad, samples=None,
                                                                            indices=None, N_MC=resolution,
                                                                            gradient=True, eps=0.005,
                                                                            border_gradients=False,
                                                                            smooth_borders=False,
                                                                            dist_threshold_ratio=dist_threshold_ratio)

        vertices, triangles = pred_mesh.vertices, pred_mesh.faces
        if world_space:
            vertices = vertices * self.dataset.scale_mats_np[0][0, 0] + self.dataset.scale_mats_np[0][:3, 3][None]

        mesh = trimesh.Trimesh(vertices, triangles)

        os.makedirs(os.path.join(self.base_exp_dir, 'udf_meshes'), exist_ok=True)
        mesh.export(
            os.path.join(self.base_exp_dir, 'udf_meshes', 'udf_res{}_step{}.ply'.format(resolution, self.iter_step)))
//  [python]
def gen_rays_from_single_image(H, W, image, intrinsic, c2w, depth=None, mask=None):
    """
    generate rays in world space, for image image
    :param H:
    :param W:
    :param intrinsics: [3,3]
    :param c2ws: [4,4]
    :return:
    """
    device = image.device
    ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                            torch.linspace(0, W - 1, W), indexing="ij")  # pytorch's meshgrid has indexing='ij'
    p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3

    # normalized ndc uv coordinates, (-1, 1)
    ndc_u = 2 * xs / (W - 1) - 1
    ndc_v = 2 * ys / (H - 1) - 1
    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)

    intrinsic_inv = torch.inverse(intrinsic)

    p = p.view(-1, 3).float().to(device)  # N_rays, 3
    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays, 3
    rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # N_rays, 3
    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays, 3
    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays, 3

    image = image.permute(1, 2, 0)
    color = image.view(-1, 3)
    depth = depth.view(-1, 1) if depth is not None else None
    mask = mask.view(-1, 1) if mask is not None else torch.ones([H * W, 1]).to(device)
    sample = {
        'rays_o': rays_o,
        'rays_v': rays_v,
        'rays_ndc_uv': rays_ndc_uv,
        'rays_color': color,
        # 'rays_depth': depth,
        'rays_mask': mask,
        'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
    }
    if depth is not None:
        sample['rays_depth'] = depth

    return sample
//  [python]
def gen_rays_from_single_image(H, W, image, intrinsic, c2w, depth=None, mask=None):
    """
    generate rays in world space, for image image
    :param H:
    :param W:
    :param intrinsics: [3,3]
    :param c2ws: [4,4]
    :return:
    """
    device = image.device
    ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                            torch.linspace(0, W - 1, W), indexing="ij")  # pytorch's meshgrid has indexing='ij'
    p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3

    # normalized ndc uv coordinates, (-1, 1)
    ndc_u = 2 * xs / (W - 1) - 1
    ndc_v = 2 * ys / (H - 1) - 1
    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)

    intrinsic_inv = torch.inverse(intrinsic)

    p = p.view(-1, 3).float().to(device)  # N_rays, 3
    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays, 3
    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_rays, 3
    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays, 3
    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays, 3

    image = image.permute(1, 2, 0)
    color = image.view(-1, 3)
    depth = depth.view(-1, 1) if depth is not None else None
    mask = mask.view(-1, 1) if mask is not None else torch.ones([H * W, 1]).to(device)
    sample = {
        'rays_o': rays_o,
        'rays_v': rays_v,
        'rays_ndc_uv': rays_ndc_uv,
        'rays_color': color,
        # 'rays_depth': depth,
        'rays_mask': mask,
        'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
    }
    if depth is not None:
        sample['rays_depth'] = depth

    return sample
//  [python]
def normalize(v: torch.Tensor) -> torch.Tensor:
    return v / torch.norm(v, dim=-1, keepdim=True)
//  [python]
def normalize(v: torch.Tensor) -> torch.Tensor:
    return v / torch.linalg.norm(v, dim=-1, keepdim=True)
//  [python]
def gen_rays_between(c2w_0, c2w_1, intrinsic, ratio, H, W, resolution_level=1):
    device = c2w_0.device

    l = resolution_level
    tx = torch.linspace(0, W - 1, W // l)
    ty = torch.linspace(0, H - 1, H // l)
    pixels_x, pixels_y = torch.meshgrid(tx, ty, indexing="ij")
    p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).to(device)  # W, H, 3

    intrinsic_inv = torch.inverse(intrinsic[:3, :3])
    p = torch.matmul(intrinsic_inv[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
    trans = c2w_0[:3, 3] * (1.0 - ratio) + c2w_1[:3, 3] * ratio

    pose_0 = c2w_0.detach().cpu().numpy()
    pose_1 = c2w_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]
    key_rots = [rot_0, rot_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 = np.linalg.inv(pose)

    c2w = torch.from_numpy(pose).to(device)
    rot = torch.from_numpy(pose[:3, :3]).cuda()
    trans = torch.from_numpy(pose[: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
    return c2w, rays_o.transpose(0, 1).contiguous().view(-1, 3), rays_v.transpose(0, 1).contiguous().view(-1, 3)
//  [python]
def gen_rays_between(c2w_0, c2w_1, intrinsic, ratio, H, W, resolution_level=1):
    device = c2w_0.device

    l = resolution_level
    tx = torch.linspace(0, W - 1, W // l)
    ty = torch.linspace(0, H - 1, H // l)
    pixels_x, pixels_y = torch.meshgrid(tx, ty, indexing="ij")
    p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).to(device)  # W, H, 3

    intrinsic_inv = torch.inverse(intrinsic[:3, :3])
    p = torch.matmul(intrinsic_inv[None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
    trans = c2w_0[:3, 3] * (1.0 - ratio) + c2w_1[:3, 3] * ratio

    pose_0 = c2w_0.detach().cpu().numpy()
    pose_1 = c2w_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]
    key_rots = [rot_0, rot_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 = np.linalg.inv(pose)

    c2w = torch.from_numpy(pose).to(device)
    rot = torch.from_numpy(pose[:3, :3]).cuda()
    trans = torch.from_numpy(pose[: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
    return c2w, rays_o.transpose(0, 1).contiguous().view(-1, 3), rays_v.transpose(0, 1).contiguous().view(-1, 3)
//  [python]
def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
        """
        Render background
        """
        batch_size, n_samples = z_vals.shape

        # Section length
        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
        mid_z_vals = z_vals + dists * 0.5

        # Section midpoints
        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3

        dis_to_center = torch.norm(pts, p=2, dim=-1, keepdim=True).clip(1.0, 1e10)
        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        density, sampled_color = nerf(pts, dirs)
        alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
        alpha = alpha.reshape(batch_size, n_samples)
        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1.0 - alpha + 1e-7], -1), -1)[:, :-1]
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            "color": color,
            "sampled_color": sampled_color,
            "alpha": alpha,
            "weights": weights,
        }
//  [python]
def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
        """
        Render background
        """
        batch_size, n_samples = z_vals.shape

        # Section length
        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
        mid_z_vals = z_vals + dists * 0.5

        # Section midpoints
        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3

        dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        density, sampled_color = nerf(pts, dirs)
        alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
        alpha = alpha.reshape(batch_size, n_samples)
        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1.0 - alpha + 1e-7], -1), -1)[:, :-1]
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            "color": color,
            "sampled_color": sampled_color,
            "alpha": alpha,
            "weights": weights,
        }
//  [python]
def tensor_div(A, factor='norm'):
    r"""divide the tensor A by a constant (the norm of A).
    """
    with torch.no_grad():
        norm = torch.norm(A)

    A_out = A / norm

    return A_out, norm
//  [python]
def tensor_div(A, factor='norm'):
    r"""divide the tensor A by a constant (the norm of A).
    """
    with torch.no_grad():
        norm = torch.linalg.norm(A)

    A_out = A / norm

    return A_out, norm
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.norm(pts, p=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        # ----------------------------------------------------------------------------------------------------------
        # Use min value of [ cos, prev_cos ]
        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
        # robust when meeting situations like below:
        #
        # SDF
        # ^
        # |\          -----x----...
        # | \        /
        # |  x      x
        # |---\----/-------------> 0 level
        # |    \  /
        # |     \/
        # |
        # ----------------------------------------------------------------------------------------------------------
        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = next_z_vals - prev_z_vals
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1.0 - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        # ----------------------------------------------------------------------------------------------------------
        # Use min value of [ cos, prev_cos ]
        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
        # robust when meeting situations like below:
        #
        # SDF
        # ^
        # |\          -----x----...
        # | \        /
        # |  x      x
        # |---\----/-------------> 0 level
        # |    \  /
        # |     \/
        # |
        # ----------------------------------------------------------------------------------------------------------
        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = next_z_vals - prev_z_vals
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1.0 - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def gmres(
    A,
    b,
    x0=None,
    M=NoOp(),
    tol=1e-10,
    maxiter: int = 100,
    check=5,
    return_iters=False,
    eps=1.0e-15,
):
    """
    Solve a linear system of equations using GMRES

    Args:
        A (torch.tensor):   black matrix
        b (torch.tensor):   black RHS

    Keyword Args:
        x0 (torch.tensor):      initial guess, defaults to zeros
        M (LinearOperator):     preconditioning operator
        tol (float):            absolute tolerance for solve
        maxiter (int):          maximum number of iterations
        check (int):            how often to check the residual
        return_iters (bool):    return iteration count along with solution
        eps (float):            small number to avoid zero vectors

    Returns:
        x (torch.tensor):   block results
        k (int, optional):  iteration count
    """
    # Should fix code at some point to treat batch dimensions right
    nbatch = A.shape[0]

    # Initial guess
    if x0 is None:
        x0 = torch.zeros_like(b)
    x = x0

    # Initial residual
    r = b - A.bmm(x.unsqueeze(-1)).squeeze(-1)

    # Basis
    V = torch.zeros((nbatch, b.shape[-1], maxiter + 1), device=A.device)
    V[..., :, 0] = r / torch.norm(r, dim=-1)[:, None]

    # Preconditioned basis
    Z = torch.zeros((nbatch, b.shape[-1], maxiter + 1), device=A.device)
    Z[..., :, 0] = M.dot(V[..., :, 0])

    # Hessenberg matrix
    H = torch.zeros((nbatch, maxiter + 1, maxiter), device=A.device)

    # Residual norm
    nR = torch.norm(r, dim=-1)

    # Check for an initial zero
    if torch.all(nR < tol):
        if return_iters:
            return x0, 0
        return x0

    # Unit vector
    e1 = torch.zeros((nbatch, maxiter + 1), device=A.device)
    e1[:, 0] = nR

    # Start iterating
    for k in range(maxiter):
        # A \cdot v
        w = A.bmm(Z[..., :, k].unsqueeze(-1)).squeeze(-1)

        # Orthogonalize
        for j in range(k + 1):
            H[..., j, k] = (
                V[..., :, j].unsqueeze(-2).bmm(w.unsqueeze(-1)).squeeze(-2).squeeze(-1)
            )
            w -= H[..., j, k].unsqueeze(-1) * V[..., :, j]

        # Next Arnoldi vector
        H[..., k + 1, k] = torch.norm(w, dim=-1) + eps
        V[..., :, k + 1] = w / H[..., k + 1, k][:, None]

        # Preconditioned
        Z[..., :, k + 1] = M.dot(V[..., :, k + 1])

        # If it's time, check the new residual value
        if (k + 1) % check == 0 or (k + 1) == maxiter:
            # Project onto our basis
            y, _, _, _ = torch.linalg.lstsq(
                H[..., : k + 1, :k], e1[..., : k + 1], rcond=None
            )
            # Calculate the value of x
            x = Z[..., :, :k].bmm(y.unsqueeze(-1)).squeeze(-1) + x0
            # Calculate the residual and the norm of the residual
            r = b - A.bmm(x.unsqueeze(-1)).squeeze(-1)
            nRc = torch.norm(r, dim=-1)
            # Check for convergence
            if torch.all(nRc < tol):
                break

    if return_iters:
        return x, k
    return x
//  [python]
def gmres(
    A,
    b,
    x0=None,
    M=NoOp(),
    tol=1e-10,
    maxiter: int = 100,
    check=5,
    return_iters=False,
    eps=1.0e-15,
):
    """
    Solve a linear system of equations using GMRES

    Args:
        A (torch.tensor):   black matrix
        b (torch.tensor):   black RHS

    Keyword Args:
        x0 (torch.tensor):      initial guess, defaults to zeros
        M (LinearOperator):     preconditioning operator
        tol (float):            absolute tolerance for solve
        maxiter (int):          maximum number of iterations
        check (int):            how often to check the residual
        return_iters (bool):    return iteration count along with solution
        eps (float):            small number to avoid zero vectors

    Returns:
        x (torch.tensor):   block results
        k (int, optional):  iteration count
    """
    # Should fix code at some point to treat batch dimensions right
    nbatch = A.shape[0]

    # Initial guess
    if x0 is None:
        x0 = torch.zeros_like(b)
    x = x0

    # Initial residual
    r = b - A.bmm(x.unsqueeze(-1)).squeeze(-1)

    # Basis
    V = torch.zeros((nbatch, b.shape[-1], maxiter + 1), device=A.device)
    V[..., :, 0] = r / torch.linalg.norm(r, dim=-1)[:, None]

    # Preconditioned basis
    Z = torch.zeros((nbatch, b.shape[-1], maxiter + 1), device=A.device)
    Z[..., :, 0] = M.dot(V[..., :, 0])

    # Hessenberg matrix
    H = torch.zeros((nbatch, maxiter + 1, maxiter), device=A.device)

    # Residual norm
    nR = torch.linalg.vector_norm(r, dim=-1)

    # Check for an initial zero
    if torch.all(nR < tol):
        if return_iters:
            return x0, 0
        return x0

    # Unit vector
    e1 = torch.zeros((nbatch, maxiter + 1), device=A.device)
    e1[:, 0] = nR

    # Start iterating
    for k in range(maxiter):
        # A \cdot v
        w = A.bmm(Z[..., :, k].unsqueeze(-1)).squeeze(-1)

        # Orthogonalize
        for j in range(k + 1):
            H[..., j, k] = (
                V[..., :, j].unsqueeze(-2).bmm(w.unsqueeze(-1)).squeeze(-2).squeeze(-1)
            )
            w -= H[..., j, k].unsqueeze(-1) * V[..., :, j]

        # Next Arnoldi vector
        H[..., k + 1, k] = torch.linalg.vector_norm(w, dim=-1) + eps
        V[..., :, k + 1] = w / H[..., k + 1, k][:, None]

        # Preconditioned
        Z[..., :, k + 1] = M.dot(V[..., :, k + 1])

        # If it's time, check the new residual value
        if (k + 1) % check == 0 or (k + 1) == maxiter:
            # Project onto our basis
            y, _, _, _ = torch.linalg.lstsq(
                H[..., : k + 1, :k], e1[..., : k + 1], rcond=None
            )
            # Calculate the value of x
            x = Z[..., :, :k].bmm(y.unsqueeze(-1)).squeeze(-1) + x0
            # Calculate the residual and the norm of the residual
            r = b - A.bmm(x.unsqueeze(-1)).squeeze(-1)
            nRc = torch.linalg.vector_norm(r, dim=-1)
            # Check for convergence
            if torch.all(nRc < tol):
                break

    if return_iters:
        return x, k
    return x
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        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(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        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.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def forward(self, A: Tensor, b: Tensor, x: Optional[Tensor]=None,
                M: Optional[Tensor]=None) -> Tensor:
        '''
        Args:
            A (Tensor): the input tensor. It is assumed to be a symmetric
                positive-definite matrix. Layout is allowed to be COO, CSR, BSR, or dense.
            b (Tensor): the tensor on the right hand side. Layout could be sparse or dense
                but is only allowed to be a type that is compatible with the layout of A.
                In other words, `A @ b` operation must be supported by the layout of A.
            x (Tensor, optional): the initial guess for the solution. Default: ``None``.
            M (Tensor, optional): the preconditioner for A. Layout is allowed to be COO,
                CSR, BSR, or dense. Default: ``None``.

        Return:
            Tensor: the solved tensor. Layout is the same as the layout of b.
        '''
        if A.ndim == b.ndim:
            b = b.squeeze(-1)
        else:
            assert A.ndim == b.ndim + 1, \
                'The number of dimensions of A and b must differ by 1'
        if x is None:
            x = torch.zeros_like(b)
        bnrm2 = torch.norm(b, dim=-1)
        if (bnrm2 == 0).all():
            return b
        atol = self.tol * bnrm2
        n = b.shape[-1]

        if self.maxiter is None:
            maxiter = n * 10
        else:
            maxiter = self.maxiter
        r = b - bmv(A, x) if x.any() else b.clone()
        rho_prev, p = None, None

        for iteration in range(maxiter):
            if (torch.norm(r, dim=-1) < atol).all():
                return x

            z = bmv(M, r) if M is not None else r
            rho_cur = vecdot(r, z)
            if iteration > 0:
                beta = rho_cur / rho_prev
                p = p * beta.unsqueeze(-1) + z
            else:  # First spin
                p = torch.empty_like(r)
                p[:] = z[:]

            q = bmv(A, p)
            alpha = rho_cur / vecdot(p, q)
            x += alpha.unsqueeze(-1)*p
            r -= alpha.unsqueeze(-1)*q
            rho_prev = rho_cur

        return x
//  [python]
def forward(self, A: Tensor, b: Tensor, x: Optional[Tensor]=None,
                M: Optional[Tensor]=None) -> Tensor:
        '''
        Args:
            A (Tensor): the input tensor. It is assumed to be a symmetric
                positive-definite matrix. Layout is allowed to be COO, CSR, BSR, or dense.
            b (Tensor): the tensor on the right hand side. Layout could be sparse or dense
                but is only allowed to be a type that is compatible with the layout of A.
                In other words, `A @ b` operation must be supported by the layout of A.
            x (Tensor, optional): the initial guess for the solution. Default: ``None``.
            M (Tensor, optional): the preconditioner for A. Layout is allowed to be COO,
                CSR, BSR, or dense. Default: ``None``.

        Return:
            Tensor: the solved tensor. Layout is the same as the layout of b.
        '''
        if A.ndim == b.ndim:
            b = b.squeeze(-1)
        else:
            assert A.ndim == b.ndim + 1, \
                'The number of dimensions of A and b must differ by 1'
        if x is None:
            x = torch.zeros_like(b)
        bnrm2 = torch.linalg.norm(b, dim=-1)
        if (bnrm2 == 0).all():
            return b
        atol = self.tol * bnrm2
        n = b.shape[-1]

        if self.maxiter is None:
            maxiter = n * 10
        else:
            maxiter = self.maxiter
        r = b - bmv(A, x) if x.any() else b.clone()
        rho_prev, p = None, None

        for iteration in range(maxiter):
            if (torch.linalg.norm(r, dim=-1) < atol).all():
                return x

            z = bmv(M, r) if M is not None else r
            rho_cur = vecdot(r, z)
            if iteration > 0:
                beta = rho_cur / rho_prev
                p = p * beta.unsqueeze(-1) + z
            else:  # First spin
                p = torch.empty_like(r)
                p[:] = z[:]

            q = bmv(A, p)
            alpha = rho_cur / vecdot(p, q)
            x += alpha.unsqueeze(-1)*p
            r -= alpha.unsqueeze(-1)*q
            rho_prev = rho_cur

        return x
//  [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, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[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 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, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[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()