_fvdb_cpp.pyi

# Copyright Contributors to the OpenVDB Project
# SPDX-License-Identifier: Apache-2.0
#
import typing
from enum import Enum
from typing import ClassVar, Optional, overload

import torch

from .types import (
    DeviceIdentifier,
    ListOfListsOfTensors,
    ListOfTensors,
    LShapeSpec,
    NumericMaxRank1,
    NumericMaxRank2,
    NumericMaxRank3,
    RShapeSpec,
    Vec3d,
    Vec3dBatch,
    Vec3dBatchOrScalar,
    Vec3dOrScalar,
    Vec3i,
    Vec3iBatch,
    Vec3iOrScalar,
)

CUTLASS: ConvPackBackend
GATHER_SCATTER: ConvPackBackend
IGEMM: ConvPackBackend
LGGS: ConvPackBackend
HALO: ConvPackBackend
DENSE: ConvPackBackend
MATMUL: ConvPackBackend

class ConvPackBackend:
    __members__: ClassVar[dict] = ...  # read-only
    CUTLASS: ClassVar[ConvPackBackend] = ...
    GATHER_SCATTER: ClassVar[ConvPackBackend] = ...
    IGEMM: ClassVar[ConvPackBackend] = ...
    LGGS: ClassVar[ConvPackBackend] = ...
    HALO: ClassVar[ConvPackBackend] = ...
    DENSE: ClassVar[ConvPackBackend] = ...
    MATMUL: ClassVar[ConvPackBackend] = ...
    __entries: ClassVar[dict] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

class GaussianSplat3d:
    class ProjectionType(Enum):
        PERSPECTIVE = ...
        ORTHOGRAPHIC = ...

    log_scales: torch.Tensor
    logit_opacities: torch.Tensor
    means: torch.Tensor
    quats: torch.Tensor
    requires_grad: bool
    sh0: torch.Tensor
    shN: torch.Tensor
    def __init__(
        self,
        means: torch.Tensor,
        quats: torch.Tensor,
        log_scales: torch.Tensor,
        logit_opacities: torch.Tensor,
        sh0: torch.Tensor,
        shN: torch.Tensor,
        accumulate_mean_2d_gradients: bool = ...,
        accumulate_max_2d_radii: bool = ...,
        detach: bool = ...,
    ) -> None: ...
    @property
    def device(self) -> torch.device: ...
    @property
    def dtype(self) -> torch.dtype: ...
    @staticmethod
    def cat(
        splats: "list[GaussianSplat3d]",
        accumulate_mean_2d_gradients: bool = False,
        accumulate_max_2d_radii: bool = False,
        detach: bool = False,
    ) -> "GaussianSplat3d": ...
    def to(self, device: torch.device, dtype: torch.dtype) -> "GaussianSplat3d": ...
    def detach(self) -> "GaussianSplat3d": ...
    def detach_in_place(self) -> None: ...
    def index_select(self, indices: torch.Tensor) -> "GaussianSplat3d": ...
    def mask_select(self, mask: torch.Tensor) -> "GaussianSplat3d": ...
    def slice_select(self, begin: int, end: int, step: int) -> "GaussianSplat3d": ...
    def index_set(self, indices: torch.Tensor, value: "GaussianSplat3d") -> None: ...
    def mask_set(self, mask: torch.Tensor, value: "GaussianSplat3d") -> None: ...
    def slice_set(self, begin: int, end: int, step: int, value: "GaussianSplat3d") -> None: ...
    @property
    def sh_degree(self) -> int: ...
    @property
    def accumulate_mean_2d_gradients(self) -> bool: ...
    @accumulate_mean_2d_gradients.setter
    def accumulate_mean_2d_gradients(self, value: bool) -> None: ...
    @property
    def accumulate_max_2d_radii(self) -> bool: ...
    @accumulate_max_2d_radii.setter
    def accumulate_max_2d_radii(self, value: bool) -> None: ...
    @staticmethod
    def from_state_dict(state_dict: dict[str, torch.Tensor]) -> GaussianSplat3d: ...
    def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: ...
    def project_gaussians_for_depths(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> ProjectedGaussianSplats: ...
    def project_gaussians_for_images(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        sh_degree_to_use: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> ProjectedGaussianSplats: ...
    def project_gaussians_for_images_and_depths(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        sh_degree_to_use: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> ProjectedGaussianSplats: ...
    def render_depths(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
        backgrounds: Optional[torch.Tensor] = ...,
    ) -> tuple[torch.Tensor, torch.Tensor]: ...
    def sparse_render_depths(
        self,
        pixels_to_render: JaggedTensor,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def render_from_projected_gaussians(
        self,
        projected_gaussians: ProjectedGaussianSplats,
        crop_width: int = ...,
        crop_height: int = ...,
        crop_origin_w: int = ...,
        crop_origin_h: int = ...,
        tile_size: int = ...,
        backgrounds: Optional[torch.Tensor] = ...,
    ) -> tuple[torch.Tensor, torch.Tensor]: ...
    def render_images(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        sh_degree_to_use: int = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
        backgrounds: Optional[torch.Tensor] = ...,
    ) -> tuple[torch.Tensor, torch.Tensor]: ...
    def render_images_from_world(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        camera_model: "CameraModel" = ...,
        distortion_coeffs: Optional[torch.Tensor] = ...,
        sh_degree_to_use: int = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
        backgrounds: Optional[torch.Tensor] = ...,
        masks: Optional[torch.Tensor] = ...,
    ) -> tuple[torch.Tensor, torch.Tensor]: ...
    def sparse_render_images(
        self,
        pixels_to_render: JaggedTensor,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        sh_degree_to_use: int = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def render_images_and_depths(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        sh_degree_to_use: int = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
        backgrounds: Optional[torch.Tensor] = ...,
    ) -> tuple[torch.Tensor, torch.Tensor]: ...
    def sparse_render_images_and_depths(
        self,
        pixels_to_render: JaggedTensor,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        sh_degree_to_use: int = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def render_num_contributing_gaussians(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> tuple[torch.Tensor, torch.Tensor]: ...
    def sparse_render_num_contributing_gaussians(
        self,
        pixels_to_render: JaggedTensor,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def render_contributing_gaussian_ids(
        self,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
        top_k_contributors: int = ...,
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def sparse_render_contributing_gaussian_ids(
        self,
        pixels_to_render: JaggedTensor,
        world_to_camera_matrices: torch.Tensor,
        projection_matrices: torch.Tensor,
        image_width: int,
        image_height: int,
        near: float,
        far: float,
        projection_type: ProjectionType = ...,
        tile_size: int = ...,
        min_radius_2d: float = ...,
        eps_2d: float = ...,
        antialias: bool = ...,
        top_k_contributors: int = ...,
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def relocate_gaussians(
        self,
        log_scales: torch.Tensor,
        logit_opacities: torch.Tensor,
        ratios: torch.Tensor,
        binomial_coeffs: torch.Tensor,
        n_max: int,
        min_opacity: float,
    ) -> tuple[torch.Tensor, torch.Tensor]: ...
    def add_noise_to_means(self, noise_scale: float, t: float = ..., k: float = ...) -> None: ...
    def reset_accumulated_gradient_state(self) -> None: ...
    def save_ply(self, filename: str, metadata: dict[str, str | int | float | torch.Tensor] | None) -> None: ...
    @staticmethod
    def from_ply(
        filename: str, device: torch.device = ...
    ) -> tuple[GaussianSplat3d, dict[str, str | int | float | torch.Tensor]]: ...
    def set_state(
        self,
        means: torch.Tensor,
        quats: torch.Tensor,
        log_scales: torch.Tensor,
        logit_opacities: torch.Tensor,
        sh0: torch.Tensor,
        shN: torch.Tensor,
    ) -> None: ...
    def state_dict(self) -> dict[str, torch.Tensor]: ...
    @property
    def accumulated_gradient_step_counts(self) -> torch.Tensor: ...
    @property
    def accumulated_max_2d_radii(self) -> torch.Tensor: ...
    @property
    def accumulated_mean_2d_gradient_norms(self) -> torch.Tensor: ...
    @property
    def num_channels(self) -> int: ...
    @property
    def num_gaussians(self) -> int: ...
    @property
    def num_sh_bases(self) -> int: ...
    @property
    def opacities(self) -> torch.Tensor: ...
    @property
    def scales(self) -> torch.Tensor: ...

class GridBatch:
    max_grids_per_batch: ClassVar[int] = ...  # read-only
    @overload
    def __init__(self, device: torch.device = ...) -> None: ...
    @overload
    def __init__(self, device: str = ...) -> None: ...
    @overload
    def __init__(self, voxel_sizes: torch.Tensor, grid_origins: torch.Tensor, device: torch.device = ...) -> None: ...
    def avg_pool(
        self,
        pool_factor: Vec3iOrScalar,
        data: JaggedTensor,
        stride: Vec3iOrScalar = 0,
        coarse_grid: GridBatch | None = None,
    ) -> tuple[JaggedTensor, GridBatch]: ...
    def bbox_at(self, bi: int) -> torch.Tensor: ...
    def clip(
        self, features: JaggedTensor, ijk_min: Vec3iBatch, ijk_max: Vec3iBatch
    ) -> tuple[JaggedTensor, GridBatch]: ...
    def clipped_grid(self, ijk_min: Vec3iBatch, ijk_max: Vec3iBatch) -> GridBatch: ...
    def coarsened_grid(self, coarsening_factor: Vec3iOrScalar) -> GridBatch: ...
    def contiguous(self) -> GridBatch: ...
    def integrate_tsdf(
        self,
        voxel_truncation_distance: float,
        projection_matrices: torch.Tensor,
        cam_to_world_matrices: torch.Tensor,
        tsdf: JaggedTensor,
        weights: JaggedTensor,
        depth_images: torch.Tensor,
        weight_images: torch.Tensor | None = None,
    ) -> tuple[GridBatch, JaggedTensor, JaggedTensor]: ...
    def integrate_tsdf_with_features(
        self,
        voxel_truncation_distance: float,
        projection_matrices: torch.Tensor,
        cam_to_world_matrices: torch.Tensor,
        tsdf: JaggedTensor,
        features: JaggedTensor,
        weights: JaggedTensor,
        depth_images: torch.Tensor,
        feature_images: torch.Tensor,
        weight_images: torch.Tensor | None = None,
    ) -> tuple[GridBatch, JaggedTensor, JaggedTensor, JaggedTensor]: ...
    def conv_grid(self, kernel_size: Vec3iOrScalar, stride: Vec3iOrScalar) -> GridBatch: ...
    def coords_in_grid(self, ijk: JaggedTensor) -> JaggedTensor: ...
    def cpu(self) -> GridBatch: ...
    def cubes_in_grid(
        self, cube_centers: JaggedTensor, cube_min: Vec3dOrScalar = 0.0, cube_max: Vec3dOrScalar = 0.0
    ) -> JaggedTensor: ...
    def cubes_intersect_grid(
        self, cube_centers: JaggedTensor, cube_min: Vec3dOrScalar = 0.0, cube_max: Vec3dOrScalar = 0.0
    ) -> JaggedTensor: ...
    def cuda(self) -> GridBatch: ...
    def cum_voxels_at(self, arg0: int) -> int: ...
    def dilated_grid(self, dilation: int) -> GridBatch: ...
    def merged_grid(self, other: GridBatch) -> GridBatch: ...
    def pruned_grid(self, mask: JaggedTensor) -> GridBatch: ...
    def inject_to(self, dst_grid: GridBatch, src: JaggedTensor, dst: JaggedTensor) -> None: ...
    def dual_bbox_at(self, arg0: int) -> torch.Tensor: ...
    def dual_grid(self, exclude_border: bool = ...) -> GridBatch: ...
    def hilbert(self, offset: torch.Tensor) -> JaggedTensor: ...
    def hilbert_zyx(self, offset: torch.Tensor) -> JaggedTensor: ...
    def morton(self, offset: torch.Tensor) -> JaggedTensor: ...
    def morton_zyx(self, offset: torch.Tensor) -> JaggedTensor: ...
    def grid_to_world(self, ijk: JaggedTensor) -> JaggedTensor: ...
    def ijk_to_index(self, ijk: JaggedTensor, cumulative: bool = False) -> JaggedTensor: ...
    def ijk_to_inv_index(self, ijk: JaggedTensor, cumulative: bool = False) -> JaggedTensor: ...
    def is_contiguous(self) -> bool: ...
    def is_same(self, other: GridBatch) -> bool: ...
    def jagged_like(self, data: torch.Tensor) -> JaggedTensor: ...
    def marching_cubes(self, field: JaggedTensor, level: float) -> tuple[JaggedTensor, JaggedTensor, JaggedTensor]: ...
    def max_pool(
        self,
        pool_factor: Vec3iOrScalar,
        data: JaggedTensor,
        stride: Vec3iOrScalar,
        coarse_grid: GridBatch | None = None,
    ) -> tuple[JaggedTensor, GridBatch]: ...
    def neighbor_indexes(self, ijk: JaggedTensor, extent: int, bitshift: int) -> JaggedTensor: ...
    def num_voxels_at(self, arg0: int) -> int: ...
    def origin_at(self, arg0: int) -> torch.Tensor: ...
    def points_in_grid(self, points: JaggedTensor) -> JaggedTensor: ...
    def ray_implicit_intersection(
        self, ray_origins: JaggedTensor, ray_directions: JaggedTensor, grid_scalars: JaggedTensor, eps: float = 0.0
    ) -> JaggedTensor: ...
    def read_from_dense_cminor(self, dense_data: torch.Tensor, dense_origins: Vec3i | None = None) -> JaggedTensor: ...
    def read_from_dense_cmajor(self, dense_data: torch.Tensor, dense_origins: Vec3i | None = None) -> JaggedTensor: ...
    def sample_bezier(self, points: JaggedTensor, voxel_data: JaggedTensor) -> JaggedTensor: ...
    def sample_bezier_with_grad(
        self, points: JaggedTensor, voxel_data: JaggedTensor
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def sample_trilinear(self, points: JaggedTensor, voxel_data: JaggedTensor) -> JaggedTensor: ...
    def sample_trilinear_with_grad(
        self, points: JaggedTensor, voxel_data: JaggedTensor
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def segments_along_rays(
        self, ray_origins: JaggedTensor, ray_directions: JaggedTensor, max_segments: int, eps: float = 0.0
    ) -> JaggedTensor: ...
    def set_from_dense_grid(
        self,
        num_grids: int,
        dense_dims: Vec3i,
        ijk_min: Vec3i = ...,
        voxel_sizes: Vec3dBatchOrScalar = ...,
        origins: Vec3dBatch = ...,
        mask: torch.Tensor | None = ...,
    ) -> None: ...
    def set_from_ijk(
        self,
        ijk: JaggedTensor,
        voxel_sizes: Vec3dBatchOrScalar = ...,
        origins: Vec3dBatch = ...,
    ) -> None: ...
    def set_from_mesh(
        self,
        mesh_vertices: JaggedTensor,
        mesh_faces: JaggedTensor,
        voxel_sizes: Vec3dBatchOrScalar = ...,
        origins: Vec3dBatch = ...,
    ) -> None: ...
    def set_from_nearest_voxels_to_points(
        self, points: JaggedTensor, voxel_sizes: Vec3dBatchOrScalar = ..., origins: Vec3dBatch = ...
    ) -> None: ...
    def set_from_points(
        self,
        points: JaggedTensor,
        voxel_sizes: Vec3dBatchOrScalar = ...,
        origins: Vec3dBatch = ...,
    ) -> None: ...
    def set_global_origin(self, origin: Vec3d) -> None: ...
    def set_global_voxel_size(self, voxel_size: Vec3dOrScalar) -> None: ...
    def sparse_conv_halo(self, input: JaggedTensor, weight: torch.Tensor, variant: int = 8) -> JaggedTensor: ...
    def sparse_conv_kernel_map(
        self, kernel_size: Vec3iOrScalar, stride: Vec3iOrScalar, target_grid: GridBatch | None = None
    ) -> tuple[SparseConvPackInfo, GridBatch]: ...
    def splat_bezier(self, points: JaggedTensor, points_data: JaggedTensor) -> JaggedTensor: ...
    def splat_trilinear(self, points: JaggedTensor, points_data: JaggedTensor) -> JaggedTensor: ...
    def refine(
        self,
        subdiv_factor: Vec3iOrScalar,
        data: JaggedTensor,
        mask: JaggedTensor | None = None,
        fine_grid: GridBatch | None = None,
    ) -> tuple[JaggedTensor, GridBatch]: ...
    def refined_grid(self, subdiv_factor: Vec3iOrScalar, mask: JaggedTensor | None = ...) -> GridBatch: ...
    @overload
    def to(self, to_device: torch.device) -> GridBatch: ...
    @overload
    def to(self, to_device: str) -> GridBatch: ...
    @overload
    def to(self, to_tensor: torch.Tensor) -> GridBatch: ...
    @overload
    def to(self, to_jtensor) -> GridBatch: ...
    @overload
    def to(self, to_grid: GridBatch) -> GridBatch: ...
    def uniform_ray_samples(
        self,
        ray_origins: JaggedTensor,
        ray_directions: JaggedTensor,
        t_min: JaggedTensor,
        t_max: JaggedTensor,
        step_size: float,
        cone_angle: float = 0.0,
        include_end_segments: bool = True,
        return_midpoints: bool = False,
        eps: float = 0.0,
    ) -> JaggedTensor: ...
    def voxel_size_at(self, arg0: int) -> torch.Tensor: ...
    def voxels_along_rays(
        self,
        ray_origins: JaggedTensor,
        ray_directions: JaggedTensor,
        max_voxels: int,
        eps: float = 0.0,
        return_ijk: bool = True,
        cumulative: bool = False,
    ) -> tuple[JaggedTensor, JaggedTensor]: ...
    def world_to_grid(self, points: JaggedTensor) -> JaggedTensor: ...
    def write_to_dense_cminor(
        self,
        sparse_data: JaggedTensor,
        min_coord: Vec3iBatch | None = ...,
        grid_size: Vec3i | None = ...,
    ) -> torch.Tensor: ...
    def write_to_dense_cmajor(
        self,
        sparse_data: JaggedTensor,
        min_coord: Vec3iBatch | None = ...,
        grid_size: Vec3i | None = ...,
    ) -> torch.Tensor: ...
    def index_int(self, arg0: int) -> GridBatch: ...
    def index_slice(self, arg0: slice) -> GridBatch: ...
    @overload
    def index_list(self, arg0: list[bool]) -> GridBatch: ...
    @overload
    def index_list(self, arg0: list[int]) -> GridBatch: ...
    def index_tensor(self, arg0: torch.Tensor) -> GridBatch: ...
    @overload
    def __getitem__(self, arg0: int) -> GridBatch: ...
    @overload
    def __getitem__(self, arg0: slice) -> GridBatch: ...
    @overload
    def __getitem__(self, arg0: list[bool]) -> GridBatch: ...
    @overload
    def __getitem__(self, arg0: list[int]) -> GridBatch: ...
    @overload
    def __getitem__(self, arg0: torch.Tensor) -> GridBatch: ...
    def __iter__(self) -> typing.Iterator[GridBatch]: ...
    def __len__(self) -> int: ...
    @property
    def address(self) -> int: ...
    @property
    def bbox(self) -> torch.Tensor: ...
    @property
    def cum_voxels(self) -> torch.Tensor: ...
    @property
    def device(self) -> torch.device: ...
    @property
    def dual_bbox(self) -> torch.Tensor: ...
    @property
    def grid_count(self) -> int: ...
    @property
    def grid_to_world_matrices(self) -> torch.Tensor: ...
    @property
    def ijk(self) -> JaggedTensor: ...
    @property
    def jidx(self) -> torch.Tensor: ...
    @property
    def joffsets(self) -> torch.Tensor: ...
    @property
    def num_bytes(self) -> torch.Tensor: ...
    @property
    def num_leaf_nodes(self) -> torch.Tensor: ...
    @property
    def num_voxels(self) -> torch.Tensor: ...
    @property
    def origins(self) -> torch.Tensor: ...
    @property
    def total_bbox(self) -> torch.Tensor: ...
    @property
    def total_bytes(self) -> int: ...
    @property
    def total_leaf_nodes(self) -> int: ...
    @property
    def total_voxels(self) -> int: ...
    @property
    def viz_edge_network(self) -> tuple[JaggedTensor, JaggedTensor]: ...
    @property
    def voxel_sizes(self) -> torch.Tensor: ...
    @property
    def world_to_grid_matrices(self) -> torch.Tensor: ...

class JaggedTensor:
    jdata: torch.Tensor
    requires_grad: bool
    @overload
    def __init__(self, tensor_list: list[list[torch.Tensor]]) -> None: ...
    @overload
    def __init__(self, tensor_list: list[torch.Tensor]) -> None: ...
    @overload
    def __init__(self, tensor: torch.Tensor) -> None: ...
    def abs(self) -> JaggedTensor: ...
    def abs_(self) -> JaggedTensor: ...
    def ceil(self) -> JaggedTensor: ...
    def ceil_(self) -> JaggedTensor: ...
    def clone(self) -> JaggedTensor: ...
    def cpu(self) -> JaggedTensor: ...
    def cuda(self) -> JaggedTensor: ...
    def detach(self) -> JaggedTensor: ...
    def double(self) -> JaggedTensor: ...
    def float(self) -> JaggedTensor: ...
    def floor(self) -> JaggedTensor: ...
    def floor_(self) -> JaggedTensor: ...
    @staticmethod
    def from_data_and_indices(arg0: torch.Tensor, arg1: torch.Tensor, arg2: int) -> JaggedTensor: ...
    @staticmethod
    def from_data_and_offsets(arg0: torch.Tensor, arg1: torch.Tensor) -> JaggedTensor: ...
    @staticmethod
    def from_data_indices_and_list_ids(
        data: torch.Tensor, indices: torch.Tensor, list_ids: torch.Tensor, num_tensors: int
    ) -> JaggedTensor: ...
    @staticmethod
    def from_data_offsets_and_list_ids(
        data: torch.Tensor, offsets: torch.Tensor, list_ids: torch.Tensor
    ) -> JaggedTensor: ...
    def int(self) -> JaggedTensor: ...
    def jagged_like(self, data: torch.Tensor) -> JaggedTensor: ...
    def jflatten(self, dim: int = ...) -> JaggedTensor: ...
    def jmax(self, dim: int = ..., keepdim: bool = ...) -> list[JaggedTensor]: ...
    def jmin(self, dim: int = ..., keepdim: bool = ...) -> list[JaggedTensor]: ...
    @overload
    def jreshape(self, lshape: list[int]) -> JaggedTensor: ...
    @overload
    def jreshape(self, lshape: list[list[int]]) -> JaggedTensor: ...
    def jreshape_as(self, other: JaggedTensor | torch.Tensor) -> JaggedTensor: ...
    def jsqueeze(self, dim: int | None = None) -> JaggedTensor: ...
    def jsum(self, dim: int = ..., keepdim: bool = ...) -> JaggedTensor: ...
    def long(self) -> JaggedTensor: ...
    def requires_grad_(self, arg0: bool) -> JaggedTensor: ...
    def rmask(self, mask: torch.Tensor) -> JaggedTensor: ...
    def round(self, decimals: int = ...) -> JaggedTensor: ...
    def round_(self, decimals: int = ...) -> JaggedTensor: ...
    def sqrt(self) -> JaggedTensor: ...
    def sqrt_(self) -> JaggedTensor: ...
    @overload
    def to(self, arg0: torch.device) -> JaggedTensor: ...
    @overload
    def to(self, arg0: str) -> JaggedTensor: ...
    @overload
    def to(self, arg0: torch.dtype) -> JaggedTensor: ...
    @overload
    def to(self, device: torch.device) -> JaggedTensor: ...
    @overload
    def to(self, device: str) -> JaggedTensor: ...
    def type(self, arg0: torch.dtype) -> JaggedTensor: ...
    def type_as(self, arg0: JaggedTensor | torch.Tensor) -> JaggedTensor: ...
    def unbind(self) -> ListOfTensors | ListOfListsOfTensors: ...
    @overload
    def __add__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __add__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __add__(self, other: int) -> JaggedTensor: ...
    @overload
    def __add__(self, other: float) -> JaggedTensor: ...
    @overload
    def __eq__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __eq__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __eq__(self, other: int) -> JaggedTensor: ...
    @overload
    def __eq__(self, other: float) -> JaggedTensor: ...
    @overload
    def __floordiv__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __floordiv__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __floordiv__(self, other: int) -> JaggedTensor: ...
    @overload
    def __floordiv__(self, other: float) -> JaggedTensor: ...
    @overload
    def __ge__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __ge__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __ge__(self, other: int) -> JaggedTensor: ...
    @overload
    def __ge__(self, other: float) -> JaggedTensor: ...
    def __getitem__(self, arg0) -> JaggedTensor: ...
    @overload
    def __gt__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __gt__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __gt__(self, other: int) -> JaggedTensor: ...
    @overload
    def __gt__(self, other: float) -> JaggedTensor: ...
    @overload
    def __iadd__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __iadd__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __iadd__(self, other: int) -> JaggedTensor: ...
    @overload
    def __iadd__(self, other: float) -> JaggedTensor: ...
    @overload
    def __ifloordiv__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __ifloordiv__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __ifloordiv__(self, other: int) -> JaggedTensor: ...
    @overload
    def __ifloordiv__(self, other: float) -> JaggedTensor: ...
    @overload
    def __imod__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __imod__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __imod__(self, other: int) -> JaggedTensor: ...
    @overload
    def __imod__(self, other: float) -> JaggedTensor: ...
    @overload
    def __imul__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __imul__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __imul__(self, other: int) -> JaggedTensor: ...
    @overload
    def __imul__(self, other: float) -> JaggedTensor: ...
    @overload
    def __ipow__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __ipow__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __ipow__(self, other: int) -> JaggedTensor: ...
    @overload
    def __ipow__(self, other: float) -> JaggedTensor: ...
    @overload
    def __isub__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __isub__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __isub__(self, other: int) -> JaggedTensor: ...
    @overload
    def __isub__(self, other: float) -> JaggedTensor: ...
    @overload
    def __itruediv__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __itruediv__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __itruediv__(self, other: int) -> JaggedTensor: ...
    @overload
    def __itruediv__(self, other: float) -> JaggedTensor: ...
    @overload
    def __le__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __le__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __le__(self, other: int) -> JaggedTensor: ...
    @overload
    def __le__(self, other: float) -> JaggedTensor: ...
    def __len__(self) -> int: ...
    @overload
    def __lt__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __lt__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __lt__(self, other: int) -> JaggedTensor: ...
    @overload
    def __lt__(self, other: float) -> JaggedTensor: ...
    @overload
    def __mod__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __mod__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __mod__(self, other: int) -> JaggedTensor: ...
    @overload
    def __mod__(self, other: float) -> JaggedTensor: ...
    @overload
    def __mul__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __mul__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __mul__(self, other: int) -> JaggedTensor: ...
    @overload
    def __mul__(self, other: float) -> JaggedTensor: ...
    @overload
    def __ne__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __ne__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __ne__(self, other: int) -> JaggedTensor: ...
    @overload
    def __ne__(self, other: float) -> JaggedTensor: ...
    @overload
    def __neg__(self) -> JaggedTensor: ...
    @overload
    def __neg__(self) -> JaggedTensor: ...
    @overload
    def __pow__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __pow__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __pow__(self, other: int) -> JaggedTensor: ...
    @overload
    def __pow__(self, other: float) -> JaggedTensor: ...
    @overload
    def __sub__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __sub__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __sub__(self, other: int) -> JaggedTensor: ...
    @overload
    def __sub__(self, other: float) -> JaggedTensor: ...
    @overload
    def __truediv__(self, other: torch.Tensor) -> JaggedTensor: ...
    @overload
    def __truediv__(self, other: JaggedTensor) -> JaggedTensor: ...
    @overload
    def __truediv__(self, other: int) -> JaggedTensor: ...
    @overload
    def __truediv__(self, other: float) -> JaggedTensor: ...
    @property
    def device(self) -> torch.device: ...
    @property
    def dtype(self) -> torch.dtype: ...
    @property
    def edim(self) -> int: ...
    @property
    def eshape(self) -> list[int]: ...
    @property
    def is_cpu(self) -> bool: ...
    @property
    def is_cuda(self) -> bool: ...
    @property
    def jidx(self) -> torch.Tensor: ...
    @property
    def jlidx(self) -> torch.Tensor: ...
    @property
    def joffsets(self) -> torch.Tensor: ...
    @property
    def ldim(self) -> int: ...
    @property
    def lshape(self) -> list[int] | list[list[int]]: ...
    @property
    def num_tensors(self) -> int: ...
    @property
    def rshape(self) -> tuple[int, ...]: ...
    def __iter__(self) -> typing.Iterator[JaggedTensor]: ...

class ProjectedGaussianSplats:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def antialias(self) -> bool: ...
    @property
    def conics(self) -> torch.Tensor: ...
    @property
    def depths(self) -> torch.Tensor: ...
    @property
    def eps_2d(self) -> float: ...
    @property
    def far_plane(self) -> float: ...
    @property
    def image_height(self) -> int: ...
    @property
    def image_width(self) -> int: ...
    @property
    def means2d(self) -> torch.Tensor: ...
    @property
    def min_radius_2d(self) -> float: ...
    @property
    def near_plane(self) -> float: ...
    @property
    def opacities(self) -> torch.Tensor: ...
    @property
    def projection_type(self) -> GaussianSplat3d.ProjectionType: ...
    @property
    def radii(self) -> torch.Tensor: ...
    @property
    def render_quantities(self) -> torch.Tensor: ...
    @property
    def sh_degree_to_use(self) -> int: ...
    @property
    def tile_gaussian_ids(self) -> torch.Tensor: ...
    @property
    def tile_offsets(self) -> torch.Tensor: ...

class SparseConvPackInfo:
    def __init__(
        self, kernel_size: Vec3iOrScalar, stride: Vec3iOrScalar, source_grid: GridBatch, target_grid: GridBatch | None
    ) -> None: ...
    def build_cutlass(self, benchmark: bool = ...) -> None: ...
    def build_gather_scatter(self, use_me: bool = ...) -> None: ...
    def build_implicit_gemm(
        self,
        sorted: bool = ...,
        split_mask_num: int = ...,
        training: bool = ...,
        split_mask_num_bwd: int = ...,
        use_tf32: bool = ...,
    ) -> None: ...
    def build_lggs(self) -> None: ...
    def cpu(self) -> SparseConvPackInfo: ...
    def cuda(self) -> SparseConvPackInfo: ...
    def sparse_conv_3d(
        self, input: JaggedTensor | torch.Tensor, weights: torch.Tensor, backend: ConvPackBackend = ...
    ) -> JaggedTensor: ...
    def sparse_transpose_conv_3d(
        self, input: JaggedTensor | torch.Tensor, weights: torch.Tensor, backend: ConvPackBackend = ...
    ) -> JaggedTensor: ...
    @overload
    def to(self, to_device: torch.device) -> SparseConvPackInfo: ...
    @overload
    def to(self, to_device: str) -> SparseConvPackInfo: ...
    @property
    def block_kernel_in_idx(self) -> torch.Tensor | None: ...
    @property
    def block_kernel_ranges(self) -> torch.Tensor | None: ...
    @property
    def block_kernel_rel_out_idx(self) -> torch.Tensor | None: ...
    @property
    def halo_index_buffer(self) -> torch.Tensor | None: ...
    @property
    def kernel_size(self) -> tuple: ...
    @property
    def neighborhood_map(self) -> torch.Tensor | None: ...
    @property
    def neighborhood_sizes(self) -> torch.Tensor | None: ...
    @property
    def out_in_map(self) -> torch.Tensor | None: ...
    @property
    def out_in_map_bwd(self) -> torch.Tensor | None: ...
    @property
    def output_index_buffer(self) -> torch.Tensor | None: ...
    @property
    def reduced_sorted_mask(self) -> torch.Tensor | None: ...
    @property
    def reorder_loc(self) -> torch.Tensor | None: ...
    @property
    def reorder_loc_bwd(self) -> torch.Tensor | None: ...
    @property
    def reorder_out_in_map(self) -> torch.Tensor | None: ...
    @property
    def reorder_out_in_map_bwd(self) -> torch.Tensor | None: ...
    @property
    def sorted_mask(self) -> torch.Tensor | None: ...
    @property
    def sorted_mask_bwd_d(self) -> torch.Tensor | None: ...
    @property
    def sorted_mask_bwd_w(self) -> torch.Tensor | None: ...
    @property
    def source_grid(self) -> GridBatch: ...
    @property
    def stride(self) -> tuple: ...
    @property
    def target_grid(self) -> GridBatch: ...
    @property
    def use_me(self) -> bool: ...
    @property
    def use_tf32(self) -> bool: ...

class GaussianSplat3dView:
    @property
    def tile_size(self) -> int: ...
    @tile_size.setter
    def tile_size(self, value: int) -> None: ...
    @property
    def min_radius_2d(self) -> float: ...
    @min_radius_2d.setter
    def min_radius_2d(self, value: float) -> None: ...
    @property
    def eps_2d(self) -> float: ...
    @eps_2d.setter
    def eps_2d(self, value: float) -> None: ...
    @property
    def antialias(self) -> bool: ...
    @antialias.setter
    def antialias(self, value: bool) -> None: ...
    @property
    def rgb_rgb_rgb_sh(self) -> bool: ...
    @rgb_rgb_rgb_sh.setter
    def rgb_rgb_rgb_sh(self, value: bool) -> None: ...
    @property
    def sh_degree_to_use(self) -> int: ...
    @sh_degree_to_use.setter
    def sh_degree_to_use(self, value: int) -> None: ...
    @property
    def sh_stride_rrr_ggg_bbb(self) -> bool: ...
    @sh_stride_rrr_ggg_bbb.setter
    def sh_stride_rrr_ggg_bbb(self, value: bool) -> None: ...

class CameraView:
    @property
    def visible(self) -> bool: ...
    @visible.setter
    def visible(self, value: bool) -> None: ...
    @property
    def axis_length(self) -> float: ...
    @axis_length.setter
    def axis_length(self, value: float) -> None: ...
    @property
    def axis_thickness(self) -> float: ...
    @axis_thickness.setter
    def axis_thickness(self, value: float) -> None: ...
    @property
    def frustum_line_width(self) -> float: ...
    @frustum_line_width.setter
    def frustum_line_width(self, value: float) -> None: ...
    @property
    def frustum_scale(self) -> float: ...
    @frustum_scale.setter
    def frustum_scale(self, value: float) -> None: ...
    @property
    def frustum_color(self) -> tuple[float, float, float]: ...
    @frustum_color.setter
    def frustum_color(self, value: tuple[float, float, float]) -> None: ...

class Viewer:
    def __init__(self, ip_address: str, port: int, device_id: int, verbose: bool) -> None: ...
    def port(self) -> int: ...
    def ip_address(self) -> str: ...
    def reset(self) -> None: ...
    def add_scene(self, scene_name: str) -> None: ...
    def remove_scene(self, scene_name: str) -> None: ...
    def remove_view(self, scene_name: str, name: str) -> None: ...
    def add_gaussian_splat_3d_view(
        self, scene_name: str, name: str, gaussian_splat_3d: GaussianSplat3d
    ) -> GaussianSplat3dView: ...
    def has_gaussian_splat_3d_view(self, name: str) -> bool: ...
    def get_gaussian_splat_3d_view(self, name: str) -> GaussianSplat3dView: ...
    def camera_orbit_center(self, scene_name: str) -> tuple[float, float, float]: ...
    def set_camera_orbit_center(self, scene_name: str, ox: float, oy: float, oz: float) -> None: ...
    def camera_orbit_radius(self, scene_name: str) -> float: ...
    def set_camera_orbit_radius(self, scene_name: str, radius: float) -> None: ...
    def camera_view_direction(self, scene_name: str) -> tuple[float, float, float]: ...
    def set_camera_view_direction(self, scene_name: str, dx: float, dy: float, dz: float) -> None: ...
    def camera_up_direction(self, scene_name: str) -> tuple[float, float, float]: ...
    def set_camera_up_direction(self, scene_name: str, ux: float, uy: float, uz: float) -> None: ...
    def camera_near(self, scene_name: str) -> float: ...
    def set_camera_near(self, scene_name: str, near: float) -> None: ...
    def camera_far(self, scene_name: str) -> float: ...
    def set_camera_far(self, scene_name: str, far: float) -> None: ...
    def camera_projection_type(self, scene_name: str) -> str: ...
    def set_camera_projection_type(self, scene_name: str, projection_type: GaussianSplat3d.ProjectionType) -> None: ...
    def add_camera_view(
        self,
        scene_name: str,
        name: str,
        camera_to_world_matrices: NumericMaxRank3,
        projection_matrices: NumericMaxRank3,
        image_sizes: NumericMaxRank2,
        frustum_near_plane: float,
        frustum_far_plane: float,
        axis_length: float,
        axis_thickness: float,
        frustum_line_width: float,
        frustum_scale: float,
        frustum_color: tuple[float, float, float],
        visible: bool,
    ) -> CameraView: ...
    def has_camera_view(self, name: str) -> bool: ...
    def get_camera_view(self, name: str) -> CameraView: ...
    def add_image(
        self,
        scene_name: str,
        name: str,
        rgba_image: NumericMaxRank1,
        width: int,
        height: int,
    ) -> None: ...
    def wait_for_interrupt(self) -> None: ...

class config:
    enable_ultra_sparse_acceleration: ClassVar[bool] = ...
    pedantic_error_checking: ClassVar[bool] = ...
    def __init__(self, *args, **kwargs) -> None: ...

def gaussian_render_jagged(
    means: JaggedTensor,
    quats: JaggedTensor,
    scales: JaggedTensor,
    opacities: JaggedTensor,
    sh_coeffs: JaggedTensor,
    viewmats: JaggedTensor,
    Ks: JaggedTensor,
    image_width: int,
    image_height: int,
    near_plane: float = ...,
    far_plane: float = ...,
    sh_degree_to_use: int = ...,
    tile_size: int = ...,
    radius_clip: float = ...,
    eps2d: float = ...,
    antialias: bool = ...,
    render_depth_channel: bool = ...,
    return_debug_info: bool = ...,
    render_depth_only: bool = ...,
    ortho: bool = ...,
    backgrounds: Optional[torch.Tensor] = ...,
) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: ...
def evaluate_spherical_harmonics(
    sh_degree: int,
    num_cameras: int,
    sh0: torch.Tensor,
    radii: torch.Tensor,
    shN: Optional[torch.Tensor] = ...,
    view_directions: Optional[torch.Tensor] = ...,
) -> torch.Tensor: ...
@overload
def jcat(grid_batches: list[GridBatch]) -> GridBatch: ...
@overload
def jcat(jagged_tensors: list[JaggedTensor | torch.Tensor], dim: int | None = ...) -> JaggedTensor: ...
def jempty(
    lshape: LShapeSpec,
    rshape: RShapeSpec | None = ...,
    dtype: torch.dtype | None = ...,
    device: DeviceIdentifier | None = ...,
    requires_grad: bool = ...,
    pin_memory: bool = ...,
) -> JaggedTensor: ...
def jones(
    lshape: LShapeSpec,
    rshape: RShapeSpec | None = ...,
    dtype: torch.dtype | None = ...,
    device: DeviceIdentifier | None = ...,
    requires_grad: bool = ...,
    pin_memory: bool = ...,
) -> JaggedTensor: ...
def jrand(
    lshape: LShapeSpec,
    rshape: RShapeSpec | None = ...,
    dtype: torch.dtype | None = ...,
    device: DeviceIdentifier | None = ...,
    requires_grad: bool = ...,
    pin_memory: bool = ...,
) -> JaggedTensor: ...
def jrandn(
    lshape: LShapeSpec,
    rshape: RShapeSpec | None = ...,
    dtype: torch.dtype | None = ...,
    device: DeviceIdentifier | None = ...,
    requires_grad: bool = ...,
    pin_memory: bool = ...,
) -> JaggedTensor: ...
def jzeros(
    lshape: LShapeSpec,
    rshape: RShapeSpec | None = ...,
    dtype: torch.dtype | None = ...,
    device: DeviceIdentifier | None = ...,
    requires_grad: bool = ...,
    pin_memory: bool = ...,
) -> JaggedTensor: ...
@overload
def load(
    path: str,
    indices: list[int],
    device: torch.device = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(
    path: str,
    indices: list[int],
    device: str = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(
    path: str,
    index: int,
    device: torch.device = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(
    path: str,
    index: int,
    device: str = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(
    path: str,
    names: list[str],
    device: torch.device = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(
    path: str,
    names: list[str],
    device: str = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(
    path: str,
    name: str,
    device: torch.device = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(
    path: str,
    name: str,
    device: str = ...,
    verbose: bool = ...,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(path: str, device: torch.device = ..., verbose: bool = ...) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load(path: str, device: str = ..., verbose: bool = ...) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def save(
    path: str,
    grid_batch: GridBatch,
    data: JaggedTensor | None = ...,
    names: list[str] = ...,
    compressed: bool = ...,
    verbose: bool = ...,
) -> None: ...
@overload
def save(
    path: str,
    grid_batch: GridBatch,
    data: JaggedTensor | None = ...,
    name: str = ...,
    compressed: bool = ...,
    verbose: bool = ...,
) -> None: ...
def morton(ijk: torch.Tensor) -> torch.Tensor: ...
def hilbert(ijk: torch.Tensor) -> torch.Tensor: ...
def scaled_dot_product_attention(
    query: JaggedTensor | torch.Tensor,
    key: JaggedTensor | torch.Tensor,
    value: JaggedTensor | torch.Tensor,
    scale: float,
) -> JaggedTensor: ...
def volume_render(
    sigmas: torch.Tensor,
    rgbs: torch.Tensor,
    deltaTs: torch.Tensor,
    ts: torch.Tensor,
    packInfo: torch.Tensor,
    transmittanceThresh: float,
) -> list[torch.Tensor]: ...

class RollingShutterType(Enum):
    NONE = ...
    VERTICAL = ...
    HORIZONTAL = ...

class CameraModel(Enum):
    PINHOLE = ...
    OPENCV_RADTAN_5 = ...
    OPENCV_RATIONAL_8 = ...
    OPENCV_RADTAN_THIN_PRISM_9 = ...
    OPENCV_THIN_PRISM_12 = ...
    ORTHOGRAPHIC = ...