jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def update_model( self ):
        dataset_train = self.build_dataset()
        self._debug_mode = False
        if self._debug_mode and os.path.exists(self._feature_cache):
            print("using feature cache")
            with open(self._feature_cache, 'rb') as f:
                train_data = torch.load(f)
        else:
            print("generating bbox features")
            train_data = self.get_features(dataset_train)
            saved = torch.save(train_data, self._feature_cache)
        train_data = torch.linalg.norm(train_data, dim=1, ord=self._norm_degree)
        
        self.remax_model = ReMax(train_data)

        self.save_final_model()
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self._compute_center_and_radius(points.join(params), device=points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.norm(points - center, dim=1).reshape(-1, 1)
        return torch.le(norm[:, None], radius).reshape(-1, 1)
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self._compute_center_and_radius(points.join(params), device=points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.linalg.norm(points - center, dim=1).reshape(-1, 1)
        return torch.le(norm[:, None], radius).reshape(-1, 1)
//  [python]
def _get_points_inside(self, points, radius):
        norm = torch.norm(points, dim=1).reshape(-1, 1)
        inside = (norm <= radius)
        index = torch.where(inside)[0]
        return points[index]
//  [python]
def _get_points_inside(self, points, radius):
        norm = torch.linalg.norm(points, dim=1).reshape(-1, 1)
        inside = (norm <= radius)
        index = torch.where(inside)[0]
        return points[index]
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self.domain._compute_center_and_radius(points.join(params), device=points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.norm(points - center, dim=1).reshape(-1, 1)
        return torch.isclose(norm[:, None], radius).reshape(-1, 1)
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self.domain._compute_center_and_radius(points.join(params), device=points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.linalg.norm(points - center, dim=1).reshape(-1, 1)
        return torch.isclose(norm[:, None], radius).reshape(-1, 1)
//  [python]
def __call__(self, data):
        if self.ord=="sum":
            data.x = data.x - data.x.min()
            data.x.div_(data.x.sum(dim=self.dim, keepdim=True).clamp_(min=1.))
        else:
            norm = torch.norm(data.x, p=self.ord, dim=self.dim, keepdim=True).clamp_(min=1.)
            data.x.div_(norm)
        return data
//  [python]
def __call__(self, data):
        if self.ord=="sum":
            data.x = data.x - data.x.min()
            data.x.div_(data.x.sum(dim=self.dim, keepdim=True).clamp_(min=1.))
        else:
            norm = torch.linalg.norm(data.x, ord=self.ord, dim=self.dim, keepdim=True).clamp_(min=1.)
            data.x.div_(norm)
        return data
//  [python]
def compute_normal_image(depth, fg_mask, camera):
    depth = depth.permute(1, 2, 0)

    height = camera.image_height
    width = camera.image_width

    homo_2d = get_homo_2d(height, width)
    K = camera.K
    K_inv = get_inverse_intrinsic(K)
    uv = np.dot(homo_2d.reshape([-1, 3]), K_inv.T).reshape([height, width, 3])
    uv = torch.from_numpy(uv).to(depth.device)
    cam_ray_dir = F.normalize(uv, p=2, dim=-1)

    pred_points = cam_ray_dir * depth

    zs = pred_points[:, :, 2]
    xs = pred_points[:, :, 0]
    ys = pred_points[:, :, 1]
    eps = 1e-10

    zy = (zs[1:, :] - zs[:-1, :]) / (ys[1:, :] - ys[:-1, :] + eps)
    zx = (zs[:, 1:] - zs[:, :-1]) / (xs[:, 1:] - xs[:, :-1] + eps)

    ny = torch.cat([-zy, torch.zeros(1, width, device=zy.device)], dim=0)
    nx = torch.cat([-zx, torch.zeros(height, 1, device=zx.device)], dim=1)
    nz = torch.ones(height, width, device=depth.device, dtype=torch.float32)
    pred_normals = torch.stack([nx, ny, nz], dim=0)

    n = torch.norm(pred_normals, dim=0, keepdim=True)
    pred_normals = pred_normals / n

    pred_normals[:, ~fg_mask] = -1
    pred_normals = ((pred_normals + 1) / 2.0).clip(0.0, 1.0)

    return pred_normals
//  [python]
def compute_normal_image(depth, fg_mask, camera):
    depth = depth.permute(1, 2, 0)

    height = camera.image_height
    width = camera.image_width

    homo_2d = get_homo_2d(height, width)
    K = camera.K
    K_inv = get_inverse_intrinsic(K)
    uv = np.dot(homo_2d.reshape([-1, 3]), K_inv.T).reshape([height, width, 3])
    uv = torch.from_numpy(uv).to(depth.device)
    cam_ray_dir = F.normalize(uv, p=2, dim=-1)

    pred_points = cam_ray_dir * depth

    zs = pred_points[:, :, 2]
    xs = pred_points[:, :, 0]
    ys = pred_points[:, :, 1]
    eps = 1e-10

    zy = (zs[1:, :] - zs[:-1, :]) / (ys[1:, :] - ys[:-1, :] + eps)
    zx = (zs[:, 1:] - zs[:, :-1]) / (xs[:, 1:] - xs[:, :-1] + eps)

    ny = torch.cat([-zy, torch.zeros(1, width, device=zy.device)], dim=0)
    nx = torch.cat([-zx, torch.zeros(height, 1, device=zx.device)], dim=1)
    nz = torch.ones(height, width, device=depth.device, dtype=torch.float32)
    pred_normals = torch.stack([nx, ny, nz], dim=0)

    n = torch.linalg.norm(pred_normals, dim=0, keepdim=True)
    pred_normals = pred_normals / n

    pred_normals[:, ~fg_mask] = -1
    pred_normals = ((pred_normals + 1) / 2.0).clip(0.0, 1.0)

    return pred_normals
//  [python]
def place_dihedral(
    a: np.ndarray,
    b: np.ndarray,
    c: np.ndarray,
    bond_angle: float,
    bond_length: float,
    torsion_angle: float,
    use_torch: bool = False,
) -> Union[np.ndarray, torch.Tensor]:
    """
    Place the point d such that the bond angle, length, and torsion angle are satisfied
    with the series a, b, c, d.
    """
    assert a.shape == b.shape == c.shape
    assert a.shape[-1] == b.shape[-1] == c.shape[-1] == 3

    if not use_torch:
        unit_vec = lambda x: x / np.linalg.norm(x, axis=-1)
        cross = lambda x, y: np.cross(x, y, axis=-1)
    else:
        ensure_tensor = (
            lambda x: torch.tensor(x, requires_grad=False).to(a.device)
            if not isinstance(x, torch.Tensor)
            else x.to(a.device)
        )
        a, b, c, bond_angle, bond_length, torsion_angle = [
            ensure_tensor(x) for x in (a, b, c, bond_angle, bond_length, torsion_angle)
        ]
        unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
        cross = lambda x, y: torch.cross(x, y, dim=-1)

    ab = b - a
    bc = unit_vec(c - b)
    n = unit_vec(cross(ab, bc))
    nbc = cross(n, bc)

    if not use_torch:
        m = np.stack([bc, nbc, n], axis=-1)
        d = np.stack(
            [
                -bond_length * np.cos(bond_angle),
                bond_length * np.cos(torsion_angle) * np.sin(bond_angle),
                bond_length * np.sin(torsion_angle) * np.sin(bond_angle),
            ],
            axis=a.ndim - 1,
        )
        d = m.dot(d)
    else:
        m = torch.stack([bc, nbc, n], dim=-1)
        d = torch.stack(
            [
                -bond_length * torch.cos(bond_angle),
                bond_length * torch.cos(torsion_angle) * torch.sin(bond_angle),
                bond_length * torch.sin(torsion_angle) * torch.sin(bond_angle),
            ],
            dim=a.ndim - 1,
        ).type(m.dtype)
        d = torch.matmul(m, d).squeeze()

    return d + c
//  [python]
def place_dihedral(
    a: np.ndarray,
    b: np.ndarray,
    c: np.ndarray,
    bond_angle: float,
    bond_length: float,
    torsion_angle: float,
    use_torch: bool = False,
) -> Union[np.ndarray, torch.Tensor]:
    """
    Place the point d such that the bond angle, length, and torsion angle are satisfied
    with the series a, b, c, d.
    """
    assert a.shape == b.shape == c.shape
    assert a.shape[-1] == b.shape[-1] == c.shape[-1] == 3

    if not use_torch:
        unit_vec = lambda x: x / np.linalg.norm(x, axis=-1)
        cross = lambda x, y: np.cross(x, y, axis=-1)
    else:
        ensure_tensor = (
            lambda x: torch.tensor(x, requires_grad=False).to(a.device)
            if not isinstance(x, torch.Tensor)
            else x.to(a.device)
        )
        a, b, c, bond_angle, bond_length, torsion_angle = [
            ensure_tensor(x) for x in (a, b, c, bond_angle, bond_length, torsion_angle)
        ]
        unit_vec = lambda x: x / torch.linalg.norm(x, dim=-1, keepdim=True)
        cross = lambda x, y: torch.linalg.cross(x, y, dim=-1)

    ab = b - a
    bc = unit_vec(c - b)
    n = unit_vec(cross(ab, bc))
    nbc = cross(n, bc)

    if not use_torch:
        m = np.stack([bc, nbc, n], axis=-1)
        d = np.stack(
            [
                -bond_length * np.cos(bond_angle),
                bond_length * np.cos(torsion_angle) * np.sin(bond_angle),
                bond_length * np.sin(torsion_angle) * np.sin(bond_angle),
            ],
            axis=a.ndim - 1,
        )
        d = m.dot(d)
    else:
        m = torch.stack([bc, nbc, n], dim=-1)
        d = torch.stack(
            [
                -bond_length * torch.cos(bond_angle),
                bond_length * torch.cos(torsion_angle) * torch.sin(bond_angle),
                bond_length * torch.sin(torsion_angle) * torch.sin(bond_angle),
            ],
            dim=a.ndim - 1,
        ).type(m.dtype)
        d = torch.matmul(m, d).squeeze()

    return d + c
//  [python]
def coord2cross(self, coord, edge_index, norm_constant=1, batch=None):

        mean = unsorted_segment_sum(coord, batch,
                                    num_segments=batch.max() + 1,
                                    aggregation_method='mean')
        row, col = edge_index
        cross = torch.cross(coord[row] - mean[batch[row]],
                            coord[col] - mean[batch[col]], dim=1)
        norm = torch.norm(cross, dim=1, keepdim=True)
        cross = cross / (norm + norm_constant)
        return cross
//  [python]
def coord2cross(self, coord, edge_index, norm_constant=1, batch=None):

        mean = unsorted_segment_sum(coord, batch,
                                    num_segments=batch.max() + 1,
                                    aggregation_method='mean')
        row, col = edge_index
        cross = torch.cross(coord[row] - mean[batch[row]],
                            coord[col] - mean[batch[col]], dim=1)
        norm = torch.linalg.norm(cross, dim=1, keepdim=True)
        cross = cross / (norm + norm_constant)
        return cross
//  [python]
def rotation_matrix(a, b):
    """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 = t.cross(a, b)
    c = t.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (t.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = t.norm(v)
    skew_sym_mat = t.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return t.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
//  [python]
def rotation_matrix(a, b):
    """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 auto_orient_and_center_poses(
    poses,
):
    """
    From nerfstudio
    https://github.com/nerfstudio-project/nerfstudio/blob/8e0c68754b2c440e2d83864fac586cddcac52dc4/nerfstudio/cameras/camera_utils.py#L515
    """
    origins = poses[..., :3, 3]
    mean_origin = torch.mean(origins, dim=0)
    translation = mean_origin
    up = torch.mean(poses[:, :3, 1], dim=0)
    up = up / torch.norm(up)
    rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
    transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
    oriented_poses = transform @ poses
    return oriented_poses, transform
//  [python]
def auto_orient_and_center_poses(
    poses,
):
    """
    From nerfstudio
    https://github.com/nerfstudio-project/nerfstudio/blob/8e0c68754b2c440e2d83864fac586cddcac52dc4/nerfstudio/cameras/camera_utils.py#L515
    """
    origins = poses[..., :3, 3]
    mean_origin = torch.mean(origins, dim=0)
    translation = mean_origin
    up = torch.mean(poses[:, :3, 1], dim=0)
    up = up / torch.linalg.norm(up)
    rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
    transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
    oriented_poses = transform @ poses
    return oriented_poses, transform
//  [python]
def forward(self):
        # make sure torch.dot(vecB,vecN) is always positive
        tmp0 = torch.dot(self.theta, self.vecB)
        tmp1 = self.theta + (torch.nn.functional.softplus(tmp0)-tmp0)*self.vecB
        vecN = tmp1 / torch.norm(tmp1)
        self.vecN = vecN.detach()

        matH = (vecN.to(self.op_list.dtype) @ self.op_list.reshape(self.num_op,-1)).reshape(self.dim_op, self.dim_op)
        if self.beta is None:
            # if self.grad_nan_to_num is not None:
            #     matH = NANGradientToNumber.apply(matH, self.grad_nan_to_num)
            # EVC = torch.linalg.eigh(matH)[1][:,-1] #take largest eigenvalue
            EVC,num_step = eigvalsh_largest_power_iteration(matH, tag_force_full=True, vec0=None, tol=self.eigen_iteration_tol)
            vecA = ((self.op_list @ EVC) @ EVC.conj()).real
        else:
            rho = numqi.manifold.symmetric_matrix_to_trace1PSD(self.beta*matH)
            vecA = (self.op_list.conj().reshape(self.num_op,-1) @ rho.reshape(-1)).real
        self.vecA = vecA.detach()
        loss = self.factor*torch.dot(vecA, vecN) / (torch.dot(vecN, self.vecB)) # 0.5* for Gell-Mann matrix nomalization
        return loss
//  [python]
def forward(self):
        # make sure torch.dot(vecB,vecN) is always positive
        tmp0 = torch.dot(self.theta, self.vecB)
        tmp1 = self.theta + (torch.nn.functional.softplus(tmp0)-tmp0)*self.vecB
        vecN = tmp1 / torch.linalg.norm(tmp1)
        self.vecN = vecN.detach()

        matH = (vecN.to(self.op_list.dtype) @ self.op_list.reshape(self.num_op,-1)).reshape(self.dim_op, self.dim_op)
        if self.beta is None:
            # if self.grad_nan_to_num is not None:
            #     matH = NANGradientToNumber.apply(matH, self.grad_nan_to_num)
            # EVC = torch.linalg.eigh(matH)[1][:,-1] #take largest eigenvalue
            EVC,num_step = eigvalsh_largest_power_iteration(matH, tag_force_full=True, vec0=None, tol=self.eigen_iteration_tol)
            vecA = ((self.op_list @ EVC) @ EVC.conj()).real
        else:
            rho = numqi.manifold.symmetric_matrix_to_trace1PSD(self.beta*matH)
            vecA = (self.op_list.conj().reshape(self.num_op,-1) @ rho.reshape(-1)).real
        self.vecA = vecA.detach()
        loss = self.factor*torch.dot(vecA, vecN) / (torch.dot(vecN, self.vecB)) # 0.5* for Gell-Mann matrix nomalization
        return loss
//  [python]
def _dunn_index_update(data: Tensor, labels: Tensor, p: float) -> Tuple[Tensor, Tensor]:
    """Update and return variables required to compute the Dunn index.

    Args:
        data: feature vectors of shape (n_samples, n_features)
        labels: cluster labels
        p: p-norm (distance metric)

    Returns:
        intercluster_distance: intercluster distances
        max_intracluster_distance: max intracluster distances

    """
    unique_labels, inverse_indices = labels.unique(return_inverse=True)
    clusters = [data[inverse_indices == label_idx] for label_idx in range(len(unique_labels))]
    centroids = [c.mean(dim=0) for c in clusters]

    intercluster_distance = torch.norm(
        torch.stack([a - b for a, b in combinations(centroids, 2)], dim=0), p=p, dim=1
    )

    max_intracluster_distance = torch.stack([
        torch.norm(ci - mu, p=p, dim=1).max() for ci, mu in zip(clusters, centroids)
    ])

    return intercluster_distance, max_intracluster_distance
//  [python]
def _dunn_index_update(data: Tensor, labels: Tensor, p: float) -> Tuple[Tensor, Tensor]:
    """Update and return variables required to compute the Dunn index.

    Args:
        data: feature vectors of shape (n_samples, n_features)
        labels: cluster labels
        p: p-norm (distance metric)

    Returns:
        intercluster_distance: intercluster distances
        max_intracluster_distance: max intracluster distances

    """
    unique_labels, inverse_indices = labels.unique(return_inverse=True)
    clusters = [data[inverse_indices == label_idx] for label_idx in range(len(unique_labels))]
    centroids = [c.mean(dim=0) for c in clusters]

    intercluster_distance = torch.linalg.norm(
        torch.stack([a - b for a, b in combinations(centroids, 2)], dim=0), ord=p, dim=1
    )

    max_intracluster_distance = torch.stack([
        torch.linalg.norm(ci - mu, ord=p, dim=1).max() for ci, mu in zip(clusters, centroids)
    ])

    return intercluster_distance, max_intracluster_distance
//  [python]
def __find_center_ind__(data: torch.Tensor, dims: tuple = (1, 2, 3)) -> int:
        """
        Calculates the center of the k-space.

        Parameters
        ----------
        data : torch.Tensor
            Input data. The shape should be (nx, ny, nc).
        dims : tuple, optional
            Dimensions to calculate the norm. Default is ``(1, 2, 3)``.

        Returns
        -------
        center_ind : int
            The center of the k-space
        """
        for dim in dims:
            data = torch.norm(data, dim=dim, keepdim=True)
        return torch.argsort(data.squeeze())[-1:]
//  [python]
def __find_center_ind__(data: torch.Tensor, dims: tuple = (1, 2, 3)) -> int:
        """
        Calculates the center of the k-space.

        Parameters
        ----------
        data : torch.Tensor
            Input data. The shape should be (nx, ny, nc).
        dims : tuple, optional
            Dimensions to calculate the norm. Default is ``(1, 2, 3)``.

        Returns
        -------
        center_ind : int
            The center of the k-space
        """
        for dim in dims:
            data = torch.linalg.norm(data, dim=dim, keepdims=True)
        return torch.argsort(data.squeeze())[-1:]
//  [python]
def test_l2(device, k):
    lhs = torch.randn(10, 16, device=device)
    rhs = torch.randn(100, 16, device=device)

    index = L2KNNIndex(rhs)
    assert index.get_emb().device == device
    assert torch.equal(index.get_emb(), rhs)

    out = index.search(lhs, k)
    assert out.score.device == device
    assert out.index.device == device
    assert out.score.size() == (10, k)
    assert out.index.size() == (10, k)

    mat = torch.norm(lhs.unsqueeze(1) - rhs.unsqueeze(0), dim=-1).pow(2)
    score, index = mat.sort(dim=-1)

    assert torch.allclose(out.score, score[:, :k])
    assert torch.equal(out.index, index[:, :k])
//  [python]
def test_l2(device, k):
    lhs = torch.randn(10, 16, device=device)
    rhs = torch.randn(100, 16, device=device)

    index = L2KNNIndex(rhs)
    assert index.get_emb().device == device
    assert torch.equal(index.get_emb(), rhs)

    out = index.search(lhs, k)
    assert out.score.device == device
    assert out.index.device == device
    assert out.score.size() == (10, k)
    assert out.index.size() == (10, k)

    mat = torch.linalg.norm(lhs.unsqueeze(1) - rhs.unsqueeze(0), dim=-1).pow(2)
    score, index = mat.sort(dim=-1)

    assert torch.allclose(out.score, score[:, :k])
    assert torch.equal(out.index, index[:, :k])
//  [python]
def generate_synthetic_data(n_samples, n_features, n_atoms, s):
    """
    Generate synthetic data for testing.
    """
    D_true = torch.randn(n_atoms, n_features)
    norms = torch.norm(D_true, dim=1) + 1e-6
    D_true /= norms[:, None]

    S_true = torch.zeros(n_samples, n_atoms)
    for i in range(n_samples):
        indices = torch.randint(0, n_atoms, (s,))
        values = torch.randn(s)
        S_true[i, indices] = values

    X = S_true @ D_true

    return X, D_true, S_true
//  [python]
def generate_synthetic_data(n_samples, n_features, n_atoms, s):
    """
    Generate synthetic data for testing.
    """
    D_true = torch.randn(n_atoms, n_features)
    norms = torch.linalg.norm(D_true, axis=1) + 1e-6
    D_true /= norms[:, None]

    S_true = torch.zeros(n_samples, n_atoms)
    for i in range(n_samples):
        indices = torch.randint(0, n_atoms, (s,))
        values = torch.randn(s)
        S_true[i, indices] = values

    X = S_true @ D_true

    return X, D_true, S_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)
    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)
    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 old_omp(X, D, s):
    """
    X: data matrix of size (n_samples, n_features)
    D: dictionary of size (n_atoms, n_features)
    s: desired sparsity level
    Returns the sparse code matrix S of size (n_samples, n_atoms)
    """
    X = torch.from_numpy(X)
    D = torch.from_numpy(D)
    n_samples, n_features = X.shape
    n_atoms, _ = D.shape

    # Normalize rows
    norms = torch.norm(D, dim=1, keepdim=True) + 1e-6
    D2 = D / norms

    residuals = X.clone()
    ids = torch.zeros((n_samples, s), dtype=int)

    for step in range(s):
        corrs = torch.abs(residuals @ D2.T)
        ids[:, step] = torch.argmax(corrs, axis=1)
        subDs = D2[ids[:, : step + 1]].mT
        m = torch.linalg.lstsq(subDs, X, rcond=None).solution
        residuals = X - (subDs @ m.unsqueeze(2)).squeeze(2)

    S = torch.zeros((n_samples, n_atoms), dtype=m.dtype)
    r = torch.arange(n_samples, dtype=int).unsqueeze(-1)
    S[r, ids] = m
    S /= norms.T
    return S.numpy()
//  [python]
def old_omp(X, D, s):
    """
    X: data matrix of size (n_samples, n_features)
    D: dictionary of size (n_atoms, n_features)
    s: desired sparsity level
    Returns the sparse code matrix S of size (n_samples, n_atoms)
    """
    X = torch.from_numpy(X)
    D = torch.from_numpy(D)
    n_samples, n_features = X.shape
    n_atoms, _ = D.shape

    # Normalize rows
    norms = torch.linalg.norm(D, axis=1, keepdims=True) + 1e-6
    D2 = D / norms

    residuals = X.clone()
    ids = torch.zeros((n_samples, s), dtype=int)

    for step in range(s):
        corrs = torch.abs(residuals @ D2.T)
        ids[:, step] = torch.argmax(corrs, axis=1)
        subDs = D2[ids[:, : step + 1]].mT
        m = torch.linalg.lstsq(subDs, X, rcond=None).solution
        residuals = X - (subDs @ m.unsqueeze(2)).squeeze(2)

    S = torch.zeros((n_samples, n_atoms), dtype=m.dtype)
    r = torch.arange(n_samples, dtype=int).unsqueeze(-1)
    S[r, ids] = m
    S /= norms.T
    return S.numpy()
//  [python]
def clip_score(text:str, image:PIL.Image, args:argparse.Namespace):
    clip_acc = 0

    with torch.no_grad():
        if use_open_clip:
            text_token = tokenizer(text).cuda()
            txt_features = model.encode_text(text_token)

            image = preprocess(image).unsqueeze(0).cuda()
            img_features = model.encode_image(image)

        else:
            inputs = tokenizer(text,return_tensors="pt")
            inputs["input_ids"] = inputs["input_ids"].cuda()
            txt_features = model.get_text_features(inputs["input_ids"])

            inputs = feature_extractor(image)
            inputs['pixel_values'] = torch.tensor(inputs['pixel_values'][0][None]).cuda()
            img_features = model.get_image_features(inputs['pixel_values'])

        img_features, txt_features = [
            x / torch.norm(x, dim=-1, keepdim=True)
            for x in [img_features, txt_features]
        ]

        clip_score = (img_features * txt_features).sum(axis=-1).cpu().numpy().item()
    
    if args.test_random_colors:
        color_gt = text.split(" ")[0]
        color_idx = color_list.index(color_gt)
        gt_idx = color_idx
    if args.test_random_textures:
        texture_gt = text.split(" ")[0]
        texture_idx = texture_list.index(texture_gt)
        gt_idx = texture_idx
    # measure the accuracy of the color prediction by evaluating the similarity between the image features and all text prompts.
    with torch.no_grad():
        similarity = (img_features * label_prompts_feats).sum(axis=-1).cpu().numpy()
        pred = np.argmax(similarity)
        if pred == gt_idx:
            clip_acc = 1
    return clip_score, clip_acc
//  [python]
def clip_score(text:str, image:PIL.Image, args:argparse.Namespace):
    clip_acc = 0

    with torch.no_grad():
        if use_open_clip:
            text_token = tokenizer(text).cuda()
            txt_features = model.encode_text(text_token)

            image = preprocess(image).unsqueeze(0).cuda()
            img_features = model.encode_image(image)

        else:
            inputs = tokenizer(text,return_tensors="pt")
            inputs["input_ids"] = inputs["input_ids"].cuda()
            txt_features = model.get_text_features(inputs["input_ids"])

            inputs = feature_extractor(image)
            inputs['pixel_values'] = torch.tensor(inputs['pixel_values'][0][None]).cuda()
            img_features = model.get_image_features(inputs['pixel_values'])

        img_features, txt_features = [
            x / torch.linalg.norm(x, axis=-1, keepdims=True)
            for x in [img_features, txt_features]
        ]

        clip_score = (img_features * txt_features).sum(axis=-1).cpu().numpy().item()
    
    if args.test_random_colors:
        color_gt = text.split(" ")[0]
        color_idx = color_list.index(color_gt)
        gt_idx = color_idx
    if args.test_random_textures:
        texture_gt = text.split(" ")[0]
        texture_idx = texture_list.index(texture_gt)
        gt_idx = texture_idx
    # measure the accuracy of the color prediction by evaluating the similarity between the image features and all text prompts.
    with torch.no_grad():
        similarity = (img_features * label_prompts_feats).sum(axis=-1).cpu().numpy()
        pred = np.argmax(similarity)
        if pred == gt_idx:
            clip_acc = 1
    return clip_score, clip_acc
//  [python]
def get_deformation(
    screw_axis,
    # Rotation params.
    with_rotation = True,
    fix_axis_vertical = False,
    # Scaling params.
    with_isotropic_scaling = False,
    min_scale = 0.5,
    max_scale = 1.5,
):
  """Get screw axis encoding of per-point rigid transformation.

  Args:
    screw_axis: [*batch, 6]. Per point, contains 3-dim rotation $r in so(3)$ and
      3 dimensions related to the translation, $v$.
    with_rotation:
    fix_axis_vertical:
    with_isotropic_scaling:
    min_scale:
    max_scale: float. Maximum scale. The scale is activated with a sigmoid
      between 0 and max_scale. Note that this affects the scale at a 0 logit.
      With max_scale=2, the zero-logit scaling is 1.

  Returns:
    er: [*batch, 3, 3]. Rotation matrix $e^r in SO(3)$ where r is the
    unnormalized axis of rotation.
    p: [*batch, 3]. Translation component of screw motion.
    with_isotropic_scaling (bool): Whether to scale points within rotation
    matrix.
  """
  identity = torch.eye(3, dtype=torch.float, device=screw_axis.device)

  if with_rotation:
    # Compute angle of rotation.
    norm = torch.norm(
        screw_axis[Ellipsis, :3], dim=-1, keepdim=True)  # [*batch, 1].
    theta = norm.unsqueeze(-1)  # [*batch, 1, 1].
    thetasq = torch.square(theta)
    thetacubed = thetasq * theta

    # Get axis of rotation and compute $[r]_x$ and $[r]_x^2$.
    axis = screw_axis[Ellipsis, :3] / norm
    tx1, tx2, tx3 = axis[Ellipsis, 0], axis[Ellipsis, 1], axis[Ellipsis, 2]  # [*batch].

    if fix_axis_vertical:
      tx1 = torch.zeros_like(tx1)
      tx3 = torch.zeros_like(tx3)

    rx = neural_fields.cross_product_matrix(tx1, tx2, tx3)  # [*batch, 3, 3].
    rxsq = torch.matmul(rx, rx)

    # Get exponentiated axis of rotation.
    # Nerfies equation 9, Rodrigues' formula.
    er = (
        identity + torch.sin(theta) / theta * rx +
        (1 - torch.cos(theta)) / thetasq * rxsq)

    # Get translation vector.
    # Nerfies equation 3 or 11.
    g = (
        identity + (1 - torch.cos(theta)) / thetasq * rx +
        (theta - torch.sin(theta)) / thetacubed * rxsq)
    p = torch.matmul(g, screw_axis[Ellipsis, 3:6, None]).squeeze(-1)
  else:
    er = identity
    p = screw_axis[Ellipsis, 3:6]

  if with_isotropic_scaling:
    scale = (
        torch.sigmoid(screw_axis[Ellipsis, 6]) * (max_scale - min_scale) + min_scale)
    new_dims = (1,) * scale.ndim
    scale_matrix = (
        identity.view(*new_dims, 3, 3) * scale.view(*scale.shape, 1, 1))
    er = torch.matmul(er, scale_matrix)

  return er, p
//  [python]
def get_deformation(
    screw_axis,
    # Rotation params.
    with_rotation = True,
    fix_axis_vertical = False,
    # Scaling params.
    with_isotropic_scaling = False,
    min_scale = 0.5,
    max_scale = 1.5,
):
  """Get screw axis encoding of per-point rigid transformation.

  Args:
    screw_axis: [*batch, 6]. Per point, contains 3-dim rotation $r in so(3)$ and
      3 dimensions related to the translation, $v$.
    with_rotation:
    fix_axis_vertical:
    with_isotropic_scaling:
    min_scale:
    max_scale: float. Maximum scale. The scale is activated with a sigmoid
      between 0 and max_scale. Note that this affects the scale at a 0 logit.
      With max_scale=2, the zero-logit scaling is 1.

  Returns:
    er: [*batch, 3, 3]. Rotation matrix $e^r in SO(3)$ where r is the
    unnormalized axis of rotation.
    p: [*batch, 3]. Translation component of screw motion.
    with_isotropic_scaling (bool): Whether to scale points within rotation
    matrix.
  """
  identity = torch.eye(3, dtype=torch.float, device=screw_axis.device)

  if with_rotation:
    # Compute angle of rotation.
    norm = torch.linalg.norm(
        screw_axis[Ellipsis, :3], dim=-1, keepdim=True)  # [*batch, 1].
    theta = norm.unsqueeze(-1)  # [*batch, 1, 1].
    thetasq = torch.square(theta)
    thetacubed = thetasq * theta

    # Get axis of rotation and compute $[r]_x$ and $[r]_x^2$.
    axis = screw_axis[Ellipsis, :3] / norm
    tx1, tx2, tx3 = axis[Ellipsis, 0], axis[Ellipsis, 1], axis[Ellipsis, 2]  # [*batch].

    if fix_axis_vertical:
      tx1 = torch.zeros_like(tx1)
      tx3 = torch.zeros_like(tx3)

    rx = neural_fields.cross_product_matrix(tx1, tx2, tx3)  # [*batch, 3, 3].
    rxsq = torch.matmul(rx, rx)

    # Get exponentiated axis of rotation.
    # Nerfies equation 9, Rodrigues' formula.
    er = (
        identity + torch.sin(theta) / theta * rx +
        (1 - torch.cos(theta)) / thetasq * rxsq)

    # Get translation vector.
    # Nerfies equation 3 or 11.
    g = (
        identity + (1 - torch.cos(theta)) / thetasq * rx +
        (theta - torch.sin(theta)) / thetacubed * rxsq)
    p = torch.matmul(g, screw_axis[Ellipsis, 3:6, None]).squeeze(-1)
  else:
    er = identity
    p = screw_axis[Ellipsis, 3:6]

  if with_isotropic_scaling:
    scale = (
        torch.sigmoid(screw_axis[Ellipsis, 6]) * (max_scale - min_scale) + min_scale)
    new_dims = (1,) * scale.ndim
    scale_matrix = (
        identity.view(*new_dims, 3, 3) * scale.view(*scale.shape, 1, 1))
    er = torch.matmul(er, scale_matrix)

  return er, p
//  [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 create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype):
        dtypes = {
            "float32": torch.float32,
            "float64": torch.float64
        }
        dtype = dtypes.get(dtype_str)

        class aten_linalg_matrix_norm(torch.nn.Module):

            def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None:
                super().__init__()
                self.ord = p
                self.dim = dim
                self.keepdim = keepdim
                self.dtype = dtype
                if self.dtype is not None:
                    self.forward = self.forward_dtype
                if out:
                    self.forward = self.forward_out
                if out_as_dtype:
                    self.forward = self.forward_prim_dtype

            def forward(self, x):
                return torch.norm(
                    x, p=self.ord, dim=self.dim, keepdim=self.keepdim
                )

            def forward_dtype(self, x):
                return torch.norm(
                    x, p=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype
                )

            def forward_prim_dtype(self, x, y):
                return torch.norm(
                    x, p=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype
                )

            def forward_out(self, x, y):
                return y, torch.norm(
                    x, p=self.ord, dim=self.dim, keepdim=self.keepdim, out=y
                )

        ref_net = None

        return aten_linalg_matrix_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_norm"
//  [python]
def create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype):
        dtypes = {
            "float32": torch.float32,
            "float64": torch.float64
        }
        dtype = dtypes.get(dtype_str)

        class aten_linalg_matrix_norm(torch.nn.Module):

            def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None:
                super().__init__()
                self.ord = p
                self.dim = dim
                self.keepdim = keepdim
                self.dtype = dtype
                if self.dtype is not None:
                    self.forward = self.forward_dtype
                if out:
                    self.forward = self.forward_out
                if out_as_dtype:
                    self.forward = self.forward_prim_dtype

            def forward(self, x):
                return torch.linalg.norm(
                    x, ord=self.ord, dim=self.dim, keepdim=self.keepdim
                )

            def forward_dtype(self, x):
                return torch.linalg.norm(
                    x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype
                )

            def forward_prim_dtype(self, x, y):
                return torch.linalg.norm(
                    x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype
                )

            def forward_out(self, x, y):
                return y, torch.linalg.norm(
                    x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, out=y
                )

        ref_net = None

        return aten_linalg_matrix_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_norm"
//  [python]
def update_one(
        self,
        img_shape,
        estimated_homography,
        true_homography,
    ):
        if estimated_homography is None:
            correctness = 0.0
            self.homography_estimation.append(correctness)
            self.errors.append(self.correctness_thresh + 1.0)
            LOG.warning("invalid matrix provided")
            return

        true_homography = true_homography.to(self.to_device).float()
        estimated_homography = estimated_homography.to(self.to_device).float()

        # img_shape = img_shape - 64

        corners = torch.tensor(
            [
                [0, 0, 1],
                [img_shape[1] - 1, 0, 1],
                [0, img_shape[0] - 1, 1],
                [img_shape[1] - 1, img_shape[0] - 1, 1],
            ],
            dtype=torch.float32,
            device=self.to_device,
        )

        # apply the true homography and the estimated homography to the corners
        real_warped_corners = torch.mm(corners, torch.transpose(true_homography, 0, 1))
        real_warped_corners = real_warped_corners[:, :2] / real_warped_corners[:, 2:]

        warped_corners = torch.mm(corners, torch.transpose(estimated_homography, 0, 1))
        warped_corners = warped_corners[:, :2] / warped_corners[:, 2:]

        mean_dist = torch.mean(
            torch.norm(real_warped_corners - warped_corners, dim=1)
        )

        # considered correct if mean distance is below a given threshold
        correctness = mean_dist <= self.correctness_thresh

        self.homography_estimation.append(correctness)
        self.errors.append(mean_dist)
//  [python]
def update_one(
        self,
        img_shape,
        estimated_homography,
        true_homography,
    ):
        if estimated_homography is None:
            correctness = 0.0
            self.homography_estimation.append(correctness)
            self.errors.append(self.correctness_thresh + 1.0)
            LOG.warning("invalid matrix provided")
            return

        true_homography = true_homography.to(self.to_device).float()
        estimated_homography = estimated_homography.to(self.to_device).float()

        # img_shape = img_shape - 64

        corners = torch.tensor(
            [
                [0, 0, 1],
                [img_shape[1] - 1, 0, 1],
                [0, img_shape[0] - 1, 1],
                [img_shape[1] - 1, img_shape[0] - 1, 1],
            ],
            dtype=torch.float32,
            device=self.to_device,
        )

        # apply the true homography and the estimated homography to the corners
        real_warped_corners = torch.mm(corners, torch.transpose(true_homography, 0, 1))
        real_warped_corners = real_warped_corners[:, :2] / real_warped_corners[:, 2:]

        warped_corners = torch.mm(corners, torch.transpose(estimated_homography, 0, 1))
        warped_corners = warped_corners[:, :2] / warped_corners[:, 2:]

        mean_dist = torch.mean(
            torch.linalg.norm(real_warped_corners - warped_corners, dim=1)
        )

        # considered correct if mean distance is below a given threshold
        correctness = mean_dist <= self.correctness_thresh

        self.homography_estimation.append(correctness)
        self.errors.append(mean_dist)
//  [python]
def norm_ft(X: torch.Tensor, ord: Optional[Union[int, float]] = None) -> torch.Tensor:
    r"""Normalize the input feature matrix with specified ``ord`` refer to pytorch's `torch.linalg.norm <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ function.

    .. note::
        The input feature matrix is expected to be a 1D vector or a 2D tensor with shape (num_samples, num_features).

    Args:
        ``X`` (``torch.Tensor``): The input feature.
        ``ord`` (``Union[int, float]``, optional): The order of the norm can be either an ``int``, ``float``. If ``ord`` is ``None``, the norm is computed with the 2-norm. Defaults to ``None``.

    Examples:
        >>> import dhg.datapipe as dd
        >>> import torch
        >>> X = torch.tensor([
                    [0.1, 0.2, 0.5],
                    [0.5, 0.2, 0.3],
                    [0.3, 0.2, 0]
                ])
        >>> dd.norm_ft(X)
        tensor([[0.1826, 0.3651, 0.9129],
                [0.8111, 0.3244, 0.4867],
                [0.8321, 0.5547, 0.0000]])
    """
    if X.dim() == 1:
        X_norm = 1 / torch.norm(X, p=ord)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    elif X.dim() == 2:
        X_norm = 1 / torch.norm(X, p=ord, dim=1, keepdim=True)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    else:
        raise ValueError(
            "The input feature matrix is expected to be a 1D verter or a 2D tensor with shape (num_samples, num_features)."
        )
//  [python]
def norm_ft(X: torch.Tensor, ord: Optional[Union[int, float]] = None) -> torch.Tensor:
    r"""Normalize the input feature matrix with specified ``ord`` refer to pytorch's `torch.linalg.norm <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ function.

    .. note::
        The input feature matrix is expected to be a 1D vector or a 2D tensor with shape (num_samples, num_features).

    Args:
        ``X`` (``torch.Tensor``): The input feature.
        ``ord`` (``Union[int, float]``, optional): The order of the norm can be either an ``int``, ``float``. If ``ord`` is ``None``, the norm is computed with the 2-norm. Defaults to ``None``.

    Examples:
        >>> import dhg.datapipe as dd
        >>> import torch
        >>> X = torch.tensor([
                    [0.1, 0.2, 0.5],
                    [0.5, 0.2, 0.3],
                    [0.3, 0.2, 0]
                ])
        >>> dd.norm_ft(X)
        tensor([[0.1826, 0.3651, 0.9129],
                [0.8111, 0.3244, 0.4867],
                [0.8321, 0.5547, 0.0000]])
    """
    if X.dim() == 1:
        X_norm = 1 / torch.linalg.norm(X, ord=ord)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    elif X.dim() == 2:
        X_norm = 1 / torch.linalg.norm(X, ord=ord, dim=1, keepdim=True)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    else:
        raise ValueError(
            "The input feature matrix is expected to be a 1D verter or a 2D tensor with shape (num_samples, num_features)."
        )
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """

        src_idx = self.src_idx[img_idx]
        src_idx = src_idx[:9]
        # src_idx = src_idx[-9:]
        idx_list = torch.cat([torch.tensor(img_idx).unsqueeze(0), src_idx], dim=0)

        poses_pair = self.pose_all[idx_list]  # [store R^-1 and C]
        intrinsics_pair = self.intrinsics_all[idx_list]
        intrinsics_inv_pair = self.intrinsics_all_inv[idx_list]
        images_gray_pair = self.images_gray[idx_list]

        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), intrinsics_pair, intrinsics_inv_pair, poses_pair, images_gray_pair