feat(pyrefly): add tensor shape type checking

[yeetypete]

Enable Pyrefly's tensor shape checking through the pyrefly-torch-stubs dev dependency. Additionally:

• Add our own shape-aware stubs for torchvision.tv_tensors.TVTensor and Image under stubs/, and declare the rank of PointCloud3D, BoundingBoxes3D, and the camera types through a TYPE_CHECKING base alias.
• Import torch submodules explicitly (torch.testing, torch.linalg, ...) since the partial stub package no longer re-exports them.
• Use tensor calls the stubs type precisely: torch.rand(()), new_empty(n, m), torch.nonzero(..., as_tuple=True), as_subclass.
• Mark the remaining stub gaps with targeted ignores linking their upstream fixes: fix(torch-stubs): declare torch.device as a context manager facebook/pyrefly#4912, fix(torch-stubs): add nn.Module.extra_repr stub facebook/pyrefly#4914, feat: let ProxyMethod resolve callable attribute targets facebook/pyrefly#4915, fix(torch-stubs): return torch.return_types named tuples facebook/pyrefly#4916.
• Enable pyrefly inlay hints in Zed and pin ruff and pyrefly to the venv binaries.
• Document the shape-checking conventions in CONTRIBUTING.md.

Assisted-by: AI

[stroxler]

Hi @yeetypete -

Excited to see this, I hope we can unblock the issues you ran into. Trying to audit for re-exports is also on my radar, I had a few local commits to improve this but got sidetracked last week.

It seems you're aware that the feature is not stable since you noted a regression in 1.3 for your uses. I actually do believe that the core type system is probably fairly stable now; at this point we have some numpy demos and pretty decent stubs for a good subset of jax + torch which gives me some confidence.

But the stubs are quite unstable and incomplete as you noted, and I suspect there are bugs lurking in what we have as well. Thanks for all the PRs, at this point every time we experiment in typing a new codebase we run into gaps and I expect that will go on for a while.

I'd love to hear how this is going for you / ideas about what we can do better. We've had good luck with a jax contributor so I know what we have can be useful for a library author, but you're in the position of being both a library author yourself and a user of torch.

If you want to discuss anything specific that's in your way or could be better, a few options:

• I can keep an eye on this conversation, but I don't want to hijack discussion of the PR to vision3d with pyrefly issues
• we have a discord and there's a tensor shapes channel
• feel free to email me (my email's on my github)

[yeetypete]

@stroxler Ah I see you've found the origin for some of the PRs, welcome! I'm experimenting with pyrefly type checking in this draft stack.

Yeah we're aware that the tensor shape typing isn't quite stable yet but I've been keeping tabs on the progress. Now that tensor subclassing (we subclass torchvision's TVTensor a lot) works this opened up the opportunity to give it a try in vision3d!

iiuc since this only impacts downstream consumers who are opting into pyrefly tensor shape checking when they consume vision3d as a dependency (and use pyrefly) the risk is low.

The overall plan is:

1. Ensure tensor shape checking remains opt-in and doesn't add any rutime dependencies for downstream consumers. pyrefly-torch-stubs needs to stay a dev dependency and all shape_extensions imports need to be behind an if TYPE_CHECKING guard.
2. Normal non-shape type checking should work for other type checkers who consume vision3d as a dependency.
3. Shape checking is compatible with Python 3.12+ / all versions of python that vision3d supports.
4. Avoid dancing around parts of the pytorch API that aren't stubbed if possible and instead upstream support for the missing stubs. As you saw here we have to do a bit of this here we the module exports.
5. Avoid adding many (any?) pyrefly ignores due to missing stubs before merging this feature if possible.

In terms of tooling that would be nice to have:

1. pyrefly check can report how many non-shape typed Tensors we have.
2. We can accidentally add incorrectly typed shape info as a return value to a function if the body is not shape type checked. Not sure if there's much to be done about this.

Wdyt? You might see some stub PRs from @simonschlaepfer and @SBetschen as well who are working on vision3d with me :)

Btw do the stubs get released at the same cadence as pyrefly or can we expect to get updates more quickly. We can also pin to github url commit sha in pyproject.toml for a bit if that works.

I'll join the discord and we can continue the discussion there :)

[stroxler]

The easy things to answer:

• Stubs are released on the same cadence as Pyrefly. We can cut a dev release pretty much anytime, so that's the way to get bleeding-edge access; we are aiming for about one 'proper' release per month but dev releases are mostly just as usable.
• At the moment, we have no support for tracking untyped tensors. That's something already on my radar, although there are questions about the best UX for this.
  • It's particularly important, I'm guessing, to find ways of identifying gradual shapes introduced implicitly because it's pretty easy to have what look like shaped operations but become gradual either through limitations of the shape transformation DSL or because the constraint solver never binds a type variable so it falls back to the gradual shape
• For correctness of shape types, our best bet is probably to use assert_shape in tests right now, which tries to do some runtime checking. I'm also wondering if we should try to build a jaxtyping-like library for runtime enforcement (possibly with some kind of global hook to opt in/out easily e.g. to turn it on for just test suites), but that's a big undertaking

The harder question is whether you can apply stubs and not impact downstream users who aren't using shape types or Pyrefly. Right now the answer is probably no, because I think another type checker is going to detect the shape types and complain. That other type checker might wind up dropping the shape types and basically doing the right thing downstream, but I'm not confident this would work. It's a pretty brittle way to fall back.

We've already had a question of how to annotate a library so that it works with other type checkers. At the moment, the only way to do it is to ship shape stubs that people can opt into (as we have with pyrefly-torch-stubs). My current thinking is that we might need either a comment or docstring segment with the Pyrefly signature or some sort of hook where we trick other type checkers into treating the types as Annotated (similar to what jaxtyping does).

This is a significant design question, and I'm not sure we'll have a good solution right away. We can potentially have a discussion with interested folks on jax and numpy (@jorenham was asking similar questions about what it would take to put shape types directly into numpy without breaking non-Pyrefly users) about what seems best / most feasible.

[yeetypete]

• For correctness of shape types, our best bet is probably to use assert_shape in tests right now, which tries to do some runtime checking.

👍 we can definitely do this.

The harder question is whether you can apply stubs and not impact downstream users who aren't using shape types or Pyrefly. Right now the answer is probably no.

My current thinking is that we might need either a comment or docstring segment with the Pyrefly signature or some sort of hook where we trick other type checkers into treating the types as Annotated (similar to what jaxtyping does).

Ah ok thanks for letting me know. as long as the above works with the pyrefly LSP I think this would be ok as a workaround.

Alternatively I guess we could do some build-time magic and autogenerate a stubs package which we release separately while stripping the pyrefly shape annotations when we build the wheel. But not sure if this is worth it or if anyone has tried something wild like this before 😂. We have no real intention of requiring the vision3d source itself to be parseable by other type checkers outside of pyrefly. In any case I think being able to annotate the shpes inline with existing type info rather than putting type annotations in a separate stub file is important to preserve.

[stroxler]

Stripping annotations might be an option if we can't figure out a cross-type-checker way of annotating libraries quickly enough (it's important but also pretty challenging to do this, so I don't know the timeline).

It might not be that hard to do the actual stripping with something like LibCST. We could potentially build a little tool for it, in the pyrefly repo or elsewhere. The big downside would be the need to publish a stub package alongside the normal package. This is basically the mirror image of what we're doing with pyrefly-torch-stubs today, it could work but it's a little annoying.

[stroxler]

Version 1.4.0dev.1 of the stubs is up and has most of your fixes in, although I think not the last one we merged. It also probably fixes most of the re-export issues you ran into, although this stack already works around those problems.

Any new issues on your side? I'm probably going to release 1.4.0dev.2 at some point this week. That should contain a rather large set of (mostly minor) torch bugfixes, plus I'm in the process of making int-arithmetic display improvements that should help a lot with convnets in particular

[stroxler]

I now have a proposal in facebook/pyrefly#5027 for a jaxtyping-esque syntax where Pyrefly types can be declared in what other type checkers see as Annotated.

I think this would give an option for using shape types without issues downstream, at the cost of heavier syntax. Whether this is better or worse than a stub-stripping tool I'm not sure, it depends how you feel about the syntax versus the complexity in shipping a separate stub package

[yeetypete]

Any new issues on your side?

Unfortunately I haven't had the chance yet to give shape checking another spin with the merged fixes. This week I'll likely have some time.

I now have a proposal in facebook/pyrefly#5027 for a jaxtyping-esque syntax where Pyrefly types can be declared in what other type checkers see as Annotated.

Nice! I'll give this a try as well. It would be nice not to have to publish a separate package with the stubs. Does the implementation mean we need to depend on shape_extensions at runtime? Or can we guard the imports behind IF_TYPE_CHECKING.

[stroxler]

Does the implementation mean we need to depend on shape_extensions at runtime? Or can we guard the imports behind IF_TYPE_CHECKING.

You'd still need shape_extensions if the annotations are evaluated. In principle if you're using only deferred or stringified annotations, it should work with if TYPE_CHECKING, just like any typing-only dependency.

FWIW I haven't actually made a production-worthy version and merged this feature yet, so you could test it but you'd have to use a pyrefly + shape_extensions built off of that draft PR.

That said, the feedback so far has been positive so I do think I will be moving forward with this in the next week or two.

[yeetypete]

FWIW I haven't actually made a production-worthy version and merged this feature yet, so you could test it but you'd have to use a pyrefly + shape_extensions built off of that draft PR.

No worries, we'll give a version built from your branch a try. @SBetschen will take over some testing here 🚀