Skip to content

feat(pyrefly): annotate tensor shapes - #182

Draft
yeetypete wants to merge 7 commits into
feat/add-tensor-shape-type-checkingfrom
feat/add-tensor-shape-annotations
Draft

yeetypete wants to merge 7 commits into
feat/add-tensor-shape-type-checkingfrom
feat/add-tensor-shape-annotations

Conversation

@yeetypete

@yeetypete yeetypete commented Sep 13, 2026

Copy link
Copy Markdown
Owner

Annotate tensor shapes across vision3d with Pyrefly's shape types, e.g. points: "Tensor[[N, int]]". Additionally:

  • Quote every shape annotation and import shape_extensions names under TYPE_CHECKING, so Python 3.12 and 3.13 never evaluate a shape type at runtime.
  • Bind integer parameters that set an output dimension with Int, e.g. voxelize(..., max_points_per_voxel: "Int[M]" = 32).
  • Make _BoxData, _MetricData, and CameraCrop generic over the dimension their fields share, and type the KITTI calibration dict.
  • Mark lines the torch stubs reject, such as Tensor.T and .item() on one-element tensors, with targeted # pyrefly: ignore comments instead of rewriting the code.
  • Teach Sphinx to render the quoted annotations: link Tensor to torch, accept dimension names, and link vision3d names in the gallery examples to our own API pages.
  • Document the conventions in CONTRIBUTING.md.

Although we introduce quite a few pyrefly: ignores here our goal should still be getting this number down to 0. We'll need to make some contributions upstream to the pyrefly torch stubs to make this happen.

Pyrefly issues encountered (LLM summary)

Stub signatures that lose or misreport shapes:

  • Tensor.T is typed Self, which is wrong for non-square matrices. mT, mH, and H are missing. Cause of the ignores in project_to_image.
  • Tensor.__iter__ yields bare Tensor, so iterating or zip-ing shaped tensors loses their shapes (accumulate_sweeps).
  • Indexing with an integer tensor, boxes[order], falls through to the bare __getitem__ overload (NMS, mAP, point sampling).
  • new_empty and the other new_* constructors take *size: int, dropping symbolic sizes (both fake kernels in _meta_registrations).
  • torch.stack over a list[Tensor[[K]]] hits the Sequence[Any] -> Tensor fallback. Only tuples are tracked.
  • torch.where returns bare Tensor in all overloads, even with a shaped condition.
  • Tensor.argsort is not defined as a method and resolves to Unknown.
  • torch.randperm is intentionally shapeless. With n: Int[N] it could return [N].
  • torch.linalg.norm with dim= returns bare.
  • Comparison operators return Any, so every mask expression and torch.where(x > 0, ...) is bare.
  • Matmul with a variadic batch, Tensor[[*Bs, 4, 4]] @ Tensor[[4, 4]], returns bare.
  • Tensor.item requires rank 0, while PyTorch accepts any one-element tensor. Source of the ignores in the ops tests.

Checker and tooling:

  • An int dimension is fully gradual: Tensor[[int, 6]] is accepted where Tensor[[8, 6]] is required. Only the rank and the other slots are checked.
  • A zero-size dimension is rejected ("dimension must be positive"), so an empty tensor such as the ground-truth _BoxData.score cannot be typed.
  • There is no shape-coverage report. pyrefly coverage treats bare Tensor as fully typed and pyrefly infer drops hints with symbolic dimensions. Coverage was measured by post-processing --report-trace.
  • A shaped return annotation over an untracked body is accepted silently. 42 of the 85 shaped returns in this PR are such fallbacks, and no rule flags them.

Assisted-by: AI

@yeetypete
yeetypete added this pull request to stack #181 September 13, 2026 00:23
@yeetypete yeetypete changed the title feat/add tensor shape annotations feat(pyrefly): annotate tensor shapes Sep 13, 2026
@yeetypete
yeetypete force-pushed the feat/add-tensor-shape-annotations branch from fb53acd to b2e186d Compare September 13, 2026 10:53
@yeetypete
yeetypete force-pushed the feat/add-tensor-shape-annotations branch from b2e186d to 74f30d2 Compare September 13, 2026 21:30
@yeetypete
yeetypete force-pushed the feat/add-tensor-shape-annotations branch from 74f30d2 to 291695b Compare September 13, 2026 21:31
@yeetypete
yeetypete force-pushed the feat/add-tensor-shape-annotations branch from 291695b to 5b2f0a4 Compare September 13, 2026 21:43
@yeetypete
yeetypete marked this pull request as draft September 14, 2026 06:12

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant