feat(pyrefly): add tensor shape type checking

.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"
       }
     }
   }

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[<error-kind>]`.
+
 ## Running tests
 
 Tests are parametrized by device via an autouse fixture in

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"]

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,

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

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

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
@@ -276,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"],
@@ -323,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

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
 
 
@@ -17,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
 
 
@@ -37,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

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)]

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:

src/vision3d/tensors/_camera.py

@@ -1,15 +1,26 @@
 """Camera-related TVTensor types."""
 
-from typing import Any, Self, override
+from typing import TYPE_CHECKING, Any, Self, override
 
 import torch
 from torch import Tensor
 from torch.utils._pytree import tree_flatten
 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.

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

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
@@ -464,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 = []

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

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:

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

stubs/torchvision/py.typed

@@ -0,0 +1 @@
+partial

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: ...