From 697690dc4ae163f64a8fdf23197d78d6401f3f8f Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 15:55:05 +0200 Subject: [PATCH 1/8] feat: add pyrefly stubs package --- pyproject.toml | 8 +++++++- uv.lock | 23 +++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c5e0792..6618140 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dev = [ # See: https://github.com/NVIDIA/cccl/issues/7896. "nvidia-cuda-cccl==13.3.4.1.2rc1", "pyrefly>=1.3.0", + "pyrefly-torch-stubs>=1.3.0", "pytest>=9.0.2", "pytest-cov>=7.1.0", "ruff>=0.15.9", @@ -60,7 +61,7 @@ branch = true [tool.pyrefly] preset = "strict" min-severity = "ignore" -search-path = ["test", "gallery", "docs/source/_ext"] +search-path = ["test", "gallery", "docs/source/_ext", "stubs"] [tool.pytest.ini_options] addopts = "--cov --cov-report=term-missing --strict-markers" @@ -73,6 +74,11 @@ filterwarnings = [ "ignore:TensorFloat32 tensor cores.*:UserWarning", ] +[tool.ruff] +# Required because the stubs use syntax newer than the Python 3.12 minimum supported +# by vision3d. Only used for development. +per-file-target-version = { "stubs/**/*.pyi" = "py314" } + [tool.ruff.lint] preview = true extend-select = ["D", "DOC", "PT", "TC"] diff --git a/uv.lock b/uv.lock index 9a17a6b..79feb7f 100644 --- a/uv.lock +++ b/uv.lock @@ -1380,6 +1380,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/6f/459f56dc345c32f8597acef13675e7e67679fb56bc3fd8f0a280eaeb8c02/pyrefly-1.3.0-py3-none-win_arm64.whl", hash = "sha256:3a3fb7c07dfb9b43d4205d8e85a83aaefc85377d0437dc8469c132836a8a7e5d", size = 14347403, upload-time = "2026-09-11T00:41:35.01Z" }, ] +[[package]] +name = "pyrefly-shape-extensions" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/47/2e/13c885edb6a5b6d0d728fb3377d92f8f05402563c736f7c60f862c67d227/pyrefly_shape_extensions-1.3.0.tar.gz", hash = "sha256:312dbbc7177a39b9990a9d5931a15104532891233bb98dc780932dd3bbb340fa", size = 9725, upload-time = "2026-09-10T23:20:50.821Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/27/4127bf31fdf5eb1c4edc934820bc964c21aae929f99d7f1c6e0798a7f1b8/pyrefly_shape_extensions-1.3.0-py3-none-any.whl", hash = "sha256:3ab1b1b83d17a7f6fa3d485999d12a13e44dbfdc93581ea9fe9ef268bb0344c4", size = 12270, upload-time = "2026-09-10T23:20:47.373Z" }, +] + +[[package]] +name = "pyrefly-torch-stubs" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyrefly-shape-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ed/f6/04755e1b93d9d42ba2ee6db02db025812522561b810879fb42b3bde57825/pyrefly_torch_stubs-1.3.0.tar.gz", hash = "sha256:b3936fd70b0fb51c048ff89e59056403a40e5e503f8b9f878435778b82bd25f8", size = 45309, upload-time = "2026-09-10T23:20:51.8Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/30/67/2376c8a019c03dacb13b8e0e183f58bd2196185c7af1ce7f8311aa015684/pyrefly_torch_stubs-1.3.0-py3-none-any.whl", hash = "sha256:94c878c38ec7f881dc7954767e3eb59da621ee7fc5afa6d3609534242c5aaa26", size = 53736, upload-time = "2026-09-10T23:20:48.666Z" }, +] + [[package]] name = "pytest" version = "9.1.1" @@ -2059,6 +2080,7 @@ dev = [ { name = "nuscenes-devkit" }, { name = "nvidia-cuda-cccl" }, { name = "pyrefly" }, + { name = "pyrefly-torch-stubs" }, { name = "pytest" }, { name = "pytest-cov" }, { name = "ruff" }, @@ -2090,6 +2112,7 @@ dev = [ { name = "nuscenes-devkit", specifier = ">=1.2.0" }, { name = "nvidia-cuda-cccl", specifier = "==13.3.4.1.2rc1" }, { name = "pyrefly", specifier = ">=1.3.0" }, + { name = "pyrefly-torch-stubs", specifier = ">=1.3.0" }, { name = "pytest", specifier = ">=9.0.2" }, { name = "pytest-cov", specifier = ">=7.1.0" }, { name = "ruff", specifier = ">=0.15.9" }, From 04d15148a46e25a829801fdb965511385808d5e2 Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 16:46:04 +0200 Subject: [PATCH 2/8] feat add tensor shape-aware stubs for untyped torchvision types --- stubs/torchvision/py.typed | 1 + stubs/torchvision/tv_tensors/_image.pyi | 19 ++++++++++++ stubs/torchvision/tv_tensors/_tv_tensor.pyi | 34 +++++++++++++++++++++ 3 files changed, 54 insertions(+) create mode 100644 stubs/torchvision/py.typed create mode 100644 stubs/torchvision/tv_tensors/_image.pyi create mode 100644 stubs/torchvision/tv_tensors/_tv_tensor.pyi diff --git a/stubs/torchvision/py.typed b/stubs/torchvision/py.typed new file mode 100644 index 0000000..b648ac9 --- /dev/null +++ b/stubs/torchvision/py.typed @@ -0,0 +1 @@ +partial diff --git a/stubs/torchvision/tv_tensors/_image.pyi b/stubs/torchvision/tv_tensors/_image.pyi new file mode 100644 index 0000000..67cceb4 --- /dev/null +++ b/stubs/torchvision/tv_tensors/_image.pyi @@ -0,0 +1,19 @@ +"""Shape-aware stub for :class:`torchvision.tv_tensors.Image`.""" + +from typing import Any, Self + +import torch +from shape_extensions import IntTuple + +from ._tv_tensor import TVTensor + +class Image[Shape: IntTuple = IntTuple](TVTensor[Shape]): + def __new__( + cls, + data: Any, + *, + dtype: torch.dtype | None = None, + device: torch.device | str | int | None = None, + requires_grad: bool | None = None, + ) -> Self: ... + def __repr__(self, *, tensor_contents: Any = None) -> str: ... diff --git a/stubs/torchvision/tv_tensors/_tv_tensor.pyi b/stubs/torchvision/tv_tensors/_tv_tensor.pyi new file mode 100644 index 0000000..49d47dd --- /dev/null +++ b/stubs/torchvision/tv_tensors/_tv_tensor.pyi @@ -0,0 +1,34 @@ +"""Shape-aware stub for :class:`torchvision.tv_tensors.TVTensor`.""" + +from collections.abc import Callable, Mapping, Sequence +from typing import Any, Self + +import torch +from shape_extensions import IntTuple +from torch import Tensor + +class TVTensor[Shape: IntTuple = IntTuple](Tensor[Shape]): + @staticmethod + def _to_tensor( + data: Any, + dtype: torch.dtype | None = None, + device: torch.device | str | int | None = None, + requires_grad: bool | None = None, + ) -> Tensor: ... + @classmethod + def _wrap_output( + cls, + output: Tensor, + args: Sequence[Any] = (), + kwargs: Mapping[str, Any] | None = None, + ) -> Tensor: ... + @classmethod + def __torch_function__( + cls, + func: Callable[..., Tensor], + types: tuple[type[Tensor], ...], + args: Sequence[Any] = (), + kwargs: Mapping[str, Any] | None = None, + ) -> Tensor: ... + def _make_repr(self, **kwargs: Any) -> str: ... + def __deepcopy__(self, memo: dict[int, Any]) -> Self: ... From dd79c5a65debebb9bb6142da2469cf059dae69c5 Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 16:54:36 +0200 Subject: [PATCH 3/8] fix: implicit torch imports --- setup.py | 2 +- src/vision3d/datasets/kitti.py | 1 + src/vision3d/datasets/nuscenes.py | 1 + src/vision3d/metrics/_nuscenes_detection_score.py | 1 + src/vision3d/ops/_meta_registrations.py | 1 + src/vision3d/transforms/_copy_paste_3d.py | 1 + src/vision3d/viz/_logging.py | 1 + test/test_metrics_nds.py | 1 + test/test_nuscenes.py | 2 ++ test/test_ops.py | 1 + test/test_ops_box3d_iou.py | 1 + test/test_ops_sweeps.py | 1 + test/test_ops_voxelize.py | 2 ++ test/test_transforms.py | 1 + test/test_transforms_v2.py | 2 ++ 15 files changed, 18 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 1d44ea5..16e9dc8 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ import os from pathlib import Path -import torch +import torch.cuda from setuptools import setup from torch.utils.cpp_extension import ( CUDA_HOME, diff --git a/src/vision3d/datasets/kitti.py b/src/vision3d/datasets/kitti.py index 297fe59..337e20d 100644 --- a/src/vision3d/datasets/kitti.py +++ b/src/vision3d/datasets/kitti.py @@ -11,6 +11,7 @@ import numpy as np import torch +import torch.linalg from torch import Tensor from torch.utils.data import Dataset from torchvision.datasets.utils import download_and_extract_archive diff --git a/src/vision3d/datasets/nuscenes.py b/src/vision3d/datasets/nuscenes.py index c5add95..1f28e50 100644 --- a/src/vision3d/datasets/nuscenes.py +++ b/src/vision3d/datasets/nuscenes.py @@ -13,6 +13,7 @@ import numpy as np import torch +import torch.linalg from torch import Tensor from torch.utils.data import Dataset from torchvision.datasets.utils import download_and_extract_archive diff --git a/src/vision3d/metrics/_nuscenes_detection_score.py b/src/vision3d/metrics/_nuscenes_detection_score.py index 2cf78fc..b942e2e 100644 --- a/src/vision3d/metrics/_nuscenes_detection_score.py +++ b/src/vision3d/metrics/_nuscenes_detection_score.py @@ -23,6 +23,7 @@ import numpy as np import torch +import torch.linalg from torch import Tensor from vision3d.metrics._types import Prediction3D, Target3D diff --git a/src/vision3d/ops/_meta_registrations.py b/src/vision3d/ops/_meta_registrations.py index c1527d6..d74f3c6 100644 --- a/src/vision3d/ops/_meta_registrations.py +++ b/src/vision3d/ops/_meta_registrations.py @@ -1,6 +1,7 @@ """Meta (fake tensor) registrations for vision3d custom ops.""" import torch +import torch.library from torch import Tensor diff --git a/src/vision3d/transforms/_copy_paste_3d.py b/src/vision3d/transforms/_copy_paste_3d.py index 2e4b2b5..fec315f 100644 --- a/src/vision3d/transforms/_copy_paste_3d.py +++ b/src/vision3d/transforms/_copy_paste_3d.py @@ -9,6 +9,7 @@ import numpy as np import torch import torch.nn.functional as F +import torch.special from PIL import Image, ImageDraw from torch import Tensor from torch.utils._pytree import tree_flatten, tree_unflatten diff --git a/src/vision3d/viz/_logging.py b/src/vision3d/viz/_logging.py index ab7ce17..9c4e0e0 100644 --- a/src/vision3d/viz/_logging.py +++ b/src/vision3d/viz/_logging.py @@ -3,6 +3,7 @@ import math import torch +import torch.linalg from torch import Tensor from vision3d.datasets import SampleInputs, SampleTargets diff --git a/test/test_metrics_nds.py b/test/test_metrics_nds.py index 466dded..302a411 100644 --- a/test/test_metrics_nds.py +++ b/test/test_metrics_nds.py @@ -12,6 +12,7 @@ import numpy as np import pytest import torch +import torch.cuda from common_utils import box_at from nuscenes.eval.common.config import config_factory from nuscenes.eval.common.data_classes import EvalBoxes diff --git a/test/test_nuscenes.py b/test/test_nuscenes.py index 2560878..8fb6818 100644 --- a/test/test_nuscenes.py +++ b/test/test_nuscenes.py @@ -13,6 +13,8 @@ import numpy as np import pytest import torch +import torch.linalg +import torch.testing from hypothesis import assume, given, settings from hypothesis import strategies as st from nuscenes.eval.detection.constants import DETECTION_NAMES as DEVKIT_DETECTION_NAMES diff --git a/test/test_ops.py b/test/test_ops.py index b506e85..e6907e1 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1,5 +1,6 @@ import pytest import torch +import torch.testing from vision3d.ops import box3d_convert from vision3d.tensors import BoundingBox3DFormat diff --git a/test/test_ops_box3d_iou.py b/test/test_ops_box3d_iou.py index 5c56964..e694fbf 100644 --- a/test/test_ops_box3d_iou.py +++ b/test/test_ops_box3d_iou.py @@ -8,6 +8,7 @@ import pytest import torch +import torch.testing from vision3d.ops import box3d_iou from vision3d.tensors import BoundingBox3DFormat diff --git a/test/test_ops_sweeps.py b/test/test_ops_sweeps.py index 93ef943..0612975 100644 --- a/test/test_ops_sweeps.py +++ b/test/test_ops_sweeps.py @@ -3,6 +3,7 @@ import numpy as np import pytest import torch +import torch.testing from hypothesis import given, settings from hypothesis import strategies as st from hypothesis.extra.numpy import arrays diff --git a/test/test_ops_voxelize.py b/test/test_ops_voxelize.py index acd8908..2a5b14d 100644 --- a/test/test_ops_voxelize.py +++ b/test/test_ops_voxelize.py @@ -5,6 +5,8 @@ import pytest import torch +import torch.cuda +import torch.testing from torch import Tensor from vision3d.ops import voxelize diff --git a/test/test_transforms.py b/test/test_transforms.py index 4e17cb4..35b0e05 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -3,6 +3,7 @@ import pytest import torch +import torch.testing from common_utils import ( check_transform, make_bounding_boxes_3d, diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index cd89dc9..4e8d1ea 100644 --- a/test/test_transforms_v2.py +++ b/test/test_transforms_v2.py @@ -4,6 +4,8 @@ import pytest import torch +import torch.linalg +import torch.testing from common_utils import make_camera_images, make_fusion_sample from torch import nn from torchvision.transforms import v2 as tv_v2 From c9824aa612d62bdc65cb612f571e8dc020fdeca1 Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 19:10:49 +0200 Subject: [PATCH 4/8] refactor: use tensor spellings the shape stubs type precisely Assisted-by: AI --- src/vision3d/ops/_meta_registrations.py | 10 +++++----- src/vision3d/ops/_nms_3d.py | 4 ++-- src/vision3d/transforms/_copy_paste_3d.py | 2 +- src/vision3d/transforms/_geometry.py | 10 +++++----- test/test_copy_paste_3d.py | 12 ++++++------ test/test_registry.py | 12 ++++++------ 6 files changed, 25 insertions(+), 25 deletions(-) diff --git a/src/vision3d/ops/_meta_registrations.py b/src/vision3d/ops/_meta_registrations.py index d74f3c6..8c07125 100644 --- a/src/vision3d/ops/_meta_registrations.py +++ b/src/vision3d/ops/_meta_registrations.py @@ -18,8 +18,8 @@ def _meta_iou_box3d(boxes1: Tensor, boxes2: Tensor) -> tuple[Tensor, Tensor]: n = boxes1.size(0) m = boxes2.size(0) # float32 regardless of input dtype (matches PyTorch3D's kernel). - vol = boxes1.new_empty((n, m), dtype=torch.float32) - iou = boxes1.new_empty((n, m), dtype=torch.float32) + vol = boxes1.new_empty(n, m, dtype=torch.float32) + iou = boxes1.new_empty(n, m, dtype=torch.float32) return vol, iou @@ -38,7 +38,7 @@ def _meta_voxelize( torch._check(len(point_cloud_range) == 6, lambda: "point_cloud_range size != 6") torch._check(len(voxel_size) == 3, lambda: "voxel_size size != 3") c = points.size(1) - voxels = points.new_empty((0, max_points_per_voxel, c)) - coords = points.new_empty((0, 3), dtype=torch.int64) - num_points = points.new_empty((0,), dtype=torch.int64) + voxels = points.new_empty(0, max_points_per_voxel, c) + coords = points.new_empty(0, 3, dtype=torch.int64) + num_points = points.new_empty(0, dtype=torch.int64) return voxels, coords, num_points diff --git a/src/vision3d/ops/_nms_3d.py b/src/vision3d/ops/_nms_3d.py index 707c3fa..a47ccbb 100644 --- a/src/vision3d/ops/_nms_3d.py +++ b/src/vision3d/ops/_nms_3d.py @@ -77,9 +77,9 @@ def batched_nms_3d( keep_mask = torch.zeros_like(scores, dtype=torch.bool) for class_id in torch.unique(idxs): - in_class = torch.where(idxs == class_id)[0] + in_class = torch.nonzero(idxs == class_id, as_tuple=True)[0] class_keep = nms_3d(boxes[in_class], scores[in_class], iou_threshold, format) keep_mask[in_class[class_keep]] = True - keep_indices = torch.where(keep_mask)[0] + keep_indices = torch.nonzero(keep_mask, as_tuple=True)[0] return keep_indices[scores[keep_indices].argsort(descending=True)] diff --git a/src/vision3d/transforms/_copy_paste_3d.py b/src/vision3d/transforms/_copy_paste_3d.py index fec315f..145c9cc 100644 --- a/src/vision3d/transforms/_copy_paste_3d.py +++ b/src/vision3d/transforms/_copy_paste_3d.py @@ -465,7 +465,7 @@ def forward(self, *inputs: Any) -> Any: for inp, tgt in zip(batch_inputs, batch_targets): self._extract_objects(inp, tgt) - if torch.rand(1).item() >= self.p: + if torch.rand(()).item() >= self.p: return tree_unflatten(flat_inputs, spec) output_inputs = [] diff --git a/src/vision3d/transforms/_geometry.py b/src/vision3d/transforms/_geometry.py index 2957a9c..705ae4f 100644 --- a/src/vision3d/transforms/_geometry.py +++ b/src/vision3d/transforms/_geometry.py @@ -122,9 +122,9 @@ def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: tx, ty, tz = self.translation_range offset = torch.tensor( [ - (torch.rand(1).item() * 2 - 1) * tx, - (torch.rand(1).item() * 2 - 1) * ty, - (torch.rand(1).item() * 2 - 1) * tz, + (torch.rand(()).item() * 2 - 1) * tx, + (torch.rand(()).item() * 2 - 1) * ty, + (torch.rand(()).item() * 2 - 1) * tz, ] ) return {"offset": offset} @@ -172,7 +172,7 @@ def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: Returns: Dict with ``"rotation_matrix"`` key containing a ``[3, 3]`` tensor. """ - angle = (torch.rand(1).item() * 2 - 1) * self.angle_range + angle = (torch.rand(()).item() * 2 - 1) * self.angle_range R = _rotation_matrix(self.axis, angle) return {"rotation_matrix": R} @@ -222,7 +222,7 @@ def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]: Dict with ``"factor"`` key. """ lo, hi = self.scale_range - factor = lo + torch.rand(1).item() * (hi - lo) + factor = lo + torch.rand(()).item() * (hi - lo) return {"factor": factor} @override diff --git a/test/test_copy_paste_3d.py b/test/test_copy_paste_3d.py index b1be498..e16443b 100644 --- a/test/test_copy_paste_3d.py +++ b/test/test_copy_paste_3d.py @@ -136,9 +136,9 @@ def _make_fusion_batch( box_data = [] all_points = [] for _j in range(num_boxes): - cx = (torch.rand(1).item() - 0.5) * 4 - cy = (torch.rand(1).item() - 0.5) * 4 - cz = torch.rand(1).item() * 10 + 5 + cx = (torch.rand(()).item() - 0.5) * 4 + cy = (torch.rand(()).item() - 0.5) * 4 + cz = torch.rand(()).item() * 10 + 5 l, w, h = 1.0, 1.0, 1.0 yaw = 0.0 box_data.append([cx, cy, cz, l, w, h, yaw]) @@ -206,9 +206,9 @@ def _make_camera_batch( for _ in range(batch_size): box_data = [] for _j in range(num_boxes): - cx = (torch.rand(1).item() - 0.5) * 4 - cy = (torch.rand(1).item() - 0.5) * 4 - cz = torch.rand(1).item() * 10 + 5 + cx = (torch.rand(()).item() - 0.5) * 4 + cy = (torch.rand(()).item() - 0.5) * 4 + cz = torch.rand(()).item() * 10 + 5 box_data.append([cx, cy, cz, 1.0, 1.0, 1.0, 0.0]) boxes = BoundingBoxes3D( diff --git a/test/test_registry.py b/test/test_registry.py index fb83c90..ba9895f 100644 --- a/test/test_registry.py +++ b/test/test_registry.py @@ -26,7 +26,7 @@ def my_kernel(inpt: torch.Tensor) -> torch.Tensor: kernel_was_called = True return inpt - t = MyTensor(torch.rand(3, 3)) + t = torch.rand(3, 3).as_subclass(MyTensor) kernel = _get_kernel(my_functional, type(t)) kernel(t) assert kernel_was_called @@ -60,7 +60,7 @@ def my_kernel(inpt: torch.Tensor) -> torch.Tensor: received_type = type(inpt) return inpt - t = CustomTVTensor(torch.rand(3, 3)) + t = torch.rand(3, 3).as_subclass(CustomTVTensor) kernel = _get_kernel(my_functional, CustomTVTensor) out = kernel(t) @@ -84,7 +84,7 @@ def my_kernel(inpt: torch.Tensor) -> torch.Tensor: received_type = type(inpt) return inpt - t = CustomTVTensor(torch.rand(3, 3)) + t = torch.rand(3, 3).as_subclass(CustomTVTensor) kernel = _get_kernel(my_functional, CustomTVTensor) kernel(t) @@ -118,7 +118,7 @@ def my_functional(inpt: torch.Tensor) -> torch.Tensor: register_kernel(my_functional, KnownTensor)(_identity_kernel) - t = UnknownTensor(torch.rand(3, 3)) + t = torch.rand(3, 3).as_subclass(UnknownTensor) kernel = _get_kernel(my_functional, type(t), allow_passthrough=True) out = kernel(t) assert out is t @@ -165,7 +165,7 @@ def parent_kernel(inpt: torch.Tensor) -> torch.Tensor: parent_called = True return inpt - child = ChildTensor(torch.rand(3, 3)) + child = torch.rand(3, 3).as_subclass(ChildTensor) kernel = _get_kernel(my_functional, type(child)) kernel(child) assert parent_called @@ -195,7 +195,7 @@ def child_kernel(inpt: torch.Tensor) -> torch.Tensor: child_called = True return inpt - child = ChildTensor(torch.rand(3, 3)) + child = torch.rand(3, 3).as_subclass(ChildTensor) kernel = _get_kernel(my_functional, type(child)) kernel(child) assert child_called From 783a0301d32b1b18eac10111138403a1f45f0fc1 Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 21:35:54 +0200 Subject: [PATCH 5/8] feat(tensors): declare tensor shapes for pyrefly --- src/vision3d/tensors/_bounding_boxes_3d.py | 12 ++++++++++-- src/vision3d/tensors/_camera.py | 19 +++++++++++++++---- src/vision3d/tensors/_point_cloud_3d.py | 11 +++++++++-- 3 files changed, 34 insertions(+), 8 deletions(-) diff --git a/src/vision3d/tensors/_bounding_boxes_3d.py b/src/vision3d/tensors/_bounding_boxes_3d.py index 48e7c8e..a150e88 100644 --- a/src/vision3d/tensors/_bounding_boxes_3d.py +++ b/src/vision3d/tensors/_bounding_boxes_3d.py @@ -1,6 +1,6 @@ from collections.abc import Mapping, Sequence from enum import Enum -from typing import Any, Self, override +from typing import TYPE_CHECKING, Any, Self, override import torch from torch import Tensor @@ -45,7 +45,15 @@ def is_rotated(format: "BoundingBox3DFormat") -> bool: ) -class BoundingBoxes3D(TVTensor): +if TYPE_CHECKING: + from shape_extensions import IntTuple + + _BoundingBoxes3DBase = TVTensor[IntTuple[int, int]] +else: + _BoundingBoxes3DBase = TVTensor + + +class BoundingBoxes3D(_BoundingBoxes3DBase): """:class:`~torch.Tensor` subclass for 3D bounding boxes with shape ``[N, K]``. Where ``N`` is the number of bounding boxes and ``K`` depends on the format: diff --git a/src/vision3d/tensors/_camera.py b/src/vision3d/tensors/_camera.py index 12184a2..04cf47d 100644 --- a/src/vision3d/tensors/_camera.py +++ b/src/vision3d/tensors/_camera.py @@ -1,6 +1,6 @@ """Camera-related TVTensor types.""" -from typing import Any, Self, override +from typing import TYPE_CHECKING, Any, Self, override import torch from torch import Tensor @@ -8,8 +8,19 @@ from torchvision import tv_tensors from torchvision.tv_tensors import TVTensor +if TYPE_CHECKING: + from shape_extensions import IntTuple -class CameraImages(tv_tensors.Image): + _CameraImagesBase = tv_tensors.Image[IntTuple[int, int, int, int]] + _CameraExtrinsicsBase = TVTensor[IntTuple[int, 4, 4]] + _CameraIntrinsicsBase = TVTensor[IntTuple[int, 3, 3]] +else: + _CameraImagesBase = tv_tensors.Image + _CameraExtrinsicsBase = TVTensor + _CameraIntrinsicsBase = TVTensor + + +class CameraImages(_CameraImagesBase): """:class:`~torchvision.tv_tensors.Image` subclass for multi-camera images with shape ``[N, C, H, W]``. Inherits from :class:`torchvision.tv_tensors.Image` so every @@ -58,7 +69,7 @@ def __repr__(self, *, tensor_contents: Any = None) -> str: return self._make_repr() -class CameraExtrinsics(TVTensor): +class CameraExtrinsics(_CameraExtrinsicsBase): """:class:`~torch.Tensor` subclass for camera extrinsic matrices with shape ``[N, 4, 4]``. Each matrix transforms a point from the dataset's *source frame* @@ -113,7 +124,7 @@ def __repr__(self, *, tensor_contents: Any = None) -> str: return self._make_repr() -class CameraIntrinsics(TVTensor): +class CameraIntrinsics(_CameraIntrinsicsBase): """:class:`~torch.Tensor` subclass for camera intrinsic matrices with shape ``[N, 3, 3]``. Each matrix maps from camera-frame 3D coordinates to pixel coordinates. diff --git a/src/vision3d/tensors/_point_cloud_3d.py b/src/vision3d/tensors/_point_cloud_3d.py index 000e4c3..c64c633 100644 --- a/src/vision3d/tensors/_point_cloud_3d.py +++ b/src/vision3d/tensors/_point_cloud_3d.py @@ -1,10 +1,17 @@ -from typing import Any, Self, override +from typing import TYPE_CHECKING, Any, Self, override import torch from torchvision.tv_tensors import TVTensor +if TYPE_CHECKING: + from shape_extensions import IntTuple -class PointCloud3D(TVTensor): + _PointCloud3DBase = TVTensor[IntTuple[int, int]] +else: + _PointCloud3DBase = TVTensor + + +class PointCloud3D(_PointCloud3DBase): """:class:`~torch.Tensor` subclass for 3D point clouds with shape ``[N, 3+C]``. The first 3 columns are ``(x, y, z)`` coordinates. Additional columns are From 8c93f869db129d409e8d76173e7da34ffdb090f7 Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 21:47:47 +0200 Subject: [PATCH 6/8] chore: mark pyrefly stub gaps with their upstream fixes Assisted-by: AI --- src/vision3d/metrics/_nuscenes_detection_score.py | 6 ++++-- src/vision3d/transforms/_transform.py | 6 ++++-- test/conftest.py | 3 ++- test/test_metrics_nds.py | 3 ++- test/test_transforms_point_cloud.py | 3 ++- 5 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/vision3d/metrics/_nuscenes_detection_score.py b/src/vision3d/metrics/_nuscenes_detection_score.py index b942e2e..b0fe85c 100644 --- a/src/vision3d/metrics/_nuscenes_detection_score.py +++ b/src/vision3d/metrics/_nuscenes_detection_score.py @@ -277,7 +277,8 @@ def update( self._check_required_fields(pred, target) # Pin internal tensors to CPU (the matching loop is sequential # Python) regardless of any ambient default-device context. - with torch.device("cpu"): + # Fixed by https://github.com/facebook/pyrefly/pull/4912. + with torch.device("cpu"): # pyrefly: ignore[bad-context-manager] pred_data = _to_box_data( pred["boxes"], pred["labels"], @@ -324,7 +325,8 @@ def compute(self) -> NuScenesDetectionScoreResult: label_aps: dict[tuple[int, float], float] = {} label_tp_errors: dict[tuple[int, str], float] = {} # Pin to CPU regardless of any ambient default-device context. - with torch.device("cpu"): + # Fixed by https://github.com/facebook/pyrefly/pull/4912. + with torch.device("cpu"): # pyrefly: ignore[bad-context-manager] for class_id in self.class_ids: period = self.orientation_periods.get(class_id, 2.0 * math.pi) tp_md: _MetricData | None = None diff --git a/src/vision3d/transforms/_transform.py b/src/vision3d/transforms/_transform.py index 0f4536d..850432c 100644 --- a/src/vision3d/transforms/_transform.py +++ b/src/vision3d/transforms/_transform.py @@ -81,8 +81,9 @@ def _needs_transform_list(self, flat_inputs: list[Any]) -> list[bool]: """ return [check_type(inpt, self._transformed_types) for inpt in flat_inputs] + # Fixed by https://github.com/facebook/pyrefly/pull/4915. @override - def forward(self, *inputs: Any) -> Any: + def forward(self, *inputs: Any) -> Any: # pyrefly: ignore[bad-override] """Apply the transform to one or more inputs (dicts, tuples, etc.). Do not override this; override :meth:`transform` instead. @@ -104,8 +105,9 @@ def forward(self, *inputs: Any) -> Any: return tree_unflatten(flat_outputs, spec) + # Fixed by https://github.com/facebook/pyrefly/pull/4914. @override - def extra_repr(self) -> str: + def extra_repr(self) -> str: # pyrefly: ignore[bad-override] """Auto-generate repr from public attributes. Returns: diff --git a/test/conftest.py b/test/conftest.py index 3e2d0d0..9a69058 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -82,7 +82,8 @@ def device(request: pytest.FixtureRequest) -> Generator[torch.device]: The active :class:`torch.device` for the test invocation. """ device = torch.device(request.param) - with device: + # Fixed by https://github.com/facebook/pyrefly/pull/4912. + with device: # pyrefly: ignore[bad-context-manager] yield device diff --git a/test/test_metrics_nds.py b/test/test_metrics_nds.py index 302a411..76e2536 100644 --- a/test/test_metrics_nds.py +++ b/test/test_metrics_nds.py @@ -450,7 +450,8 @@ def test_compute_under_cuda_default_device() -> None: metric = NuScenesDetectionScore( class_ids=[0], tp_metrics=("trans_err", "scale_err", "orient_err") ) - with torch.device("cuda"): + # Fixed by https://github.com/facebook/pyrefly/pull/4912. + with torch.device("cuda"): # pyrefly: ignore[bad-context-manager] pred, tgt = _perfect_car_frame() metric.update([pred], [tgt]) out = metric.compute() diff --git a/test/test_transforms_point_cloud.py b/test/test_transforms_point_cloud.py index f6d5c52..b54cba6 100644 --- a/test/test_transforms_point_cloud.py +++ b/test/test_transforms_point_cloud.py @@ -31,7 +31,8 @@ def test_output_is_permutation(self) -> None: perm = torch.randperm(50) out = shuffle_points_point_cloud(points, perm=perm) assert out.shape == points.shape - assert torch.allclose(out.sort(dim=0).values, points.sort(dim=0).values) + # Fixed by https://github.com/facebook/pyrefly/pull/4916. + assert torch.allclose(out.sort(dim=0).values, points.sort(dim=0).values) # pyrefly: ignore[missing-attribute] def test_does_not_modify_input(self) -> None: points = torch.randn(50, 4) From 00e542334f80a0494f929dee380ccdc8c07f9f45 Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 22:13:01 +0200 Subject: [PATCH 7/8] chore: update pyrefly and ruff configuration in zed settings --- .zed/settings.json | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/.zed/settings.json b/.zed/settings.json index 474917c..d4b5953 100644 --- a/.zed/settings.json +++ b/.zed/settings.json @@ -10,6 +10,14 @@ "source.organizeImports.ruff": true }, "format_on_save": "on", + "formatter": { + "language_server": { + "name": "ruff" + } + }, + "inlay_hints": { + "enabled": true + }, "language_servers": [ "pyrefly", "!basedpyright", @@ -25,22 +33,12 @@ "binary": { "arguments": ["lsp"], "path": ".venv/bin/pyrefly" - }, - "settings": { - "python": { - "pyrefly": { - "displayTypeErrors": "force-on" - } - } } }, "ruff": { - "initialization_options": { - "settings": { - "lint": { - "enable": true - } - } + "binary": { + "arguments": ["server"], + "path": ".venv/bin/ruff" } } } From 9cc60891a101ff472af6cfc113ecaa99b38c1345 Mon Sep 17 00:00:00 2001 From: Peter Siegel Date: Sat, 12 Sep 2026 22:33:31 +0200 Subject: [PATCH 8/8] doc: update contributing guide tensor shape checking info --- CONTRIBUTING.md | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 86c7759..2c2f53d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -83,6 +83,36 @@ uv run ruff format --check # check-only; fails if formatting is off uv run pyrefly check # type check ``` +### Tensor shapes + +vision3d is configured for Pyrefly's +[tensor shape checking](https://pyrefly.org/en/docs/tensor-shapes/). The +`pyrefly-torch-stubs` dev dependency replaces PyTorch's type stubs with +shape-aware ones. Pyrefly enables shape tracking automatically once it +resolves them. The `stubs/` directory holds our own shape-aware overlays for +third-party types that ship without shape information, so that the tensor +types vision3d builds on them are shape-aware as well. + +Shape types exist only for the type checker. Python 3.12 and 3.13 (which vision3d +supports) evaluate annotations at import time, so a shape type must never be +evaluated at runtime (see Pyrefly's +[runtime considerations](https://pyrefly.org/en/docs/tensor-shapes-setup/#imports-and-runtime-considerations)): + +- Quote shape annotations in signatures, `x: "Tensor[[N, 3]]"`, import + `IntTuple`, `IntVar`, and `Int` from `shape_extensions` under + `TYPE_CHECKING`, and bind shape variables as + [PEP 695](https://peps.python.org/pep-0695/) type parameters, + `def f[N: IntVar](...)`, which are evaluated lazily. Do not use + `from __future__ import annotations`: it turns `TypedDict` fields into + strings, which breaks `Required`/`NotRequired` at runtime (see the + [`TypedDict` note](https://docs.python.org/3/library/typing.html#typing.TypedDict.__optional_keys__)). +- Give a class its shape through a `TYPE_CHECKING` base alias with a plain + fallback in the `else` branch. Base classes are evaluated at runtime. + +Where the stubs are stricter than PyTorch or lack a definition, prefer an +equivalent call the stubs type precisely. Only if none exists, mark the line +with a targeted `# pyrefly: ignore[]`. + ## Running tests Tests are parametrized by device via an autouse fixture in