10000 Pytorch Typing, for Tensor type annotations · Issue #73359 · pytorch/pytorch · GitHub
[go: up one dir, main page]

Skip to content

Pytorch Typing, for Tensor type annotations #73359

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JulesGM opened this issue Feb 24, 2 8000 022 · 3 comments
Open

Pytorch Typing, for Tensor type annotations #73359

JulesGM opened this issue Feb 24, 2022 · 3 comments
Labels
module: typing Related to mypy type annotations triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Comments

@JulesGM
Copy link
JulesGM commented Feb 24, 2022

🚀 The feature, motivation and pitch

Doing something like https://numpy.org/devdocs/reference/typing.html#numpy.typing.NDArray would be cool, to allow people to see read the dtype of tensors in the signature, and possibly something like Beartype to do runtime type checking to make sure that the tensors are of the correct dtype, would be great. Number of dimensions could also be informative and would be runtime testable.

Allowing to name the dimensions (eg "batch_size" x "sequence_length" x "hidden_state_size") would also be great.
Some inspiration is https://github.com/patrick-kidger/torchtyping
In-depth discussion beartype/beartype#96

Alternatives

No response

Additional context

https://numpy.org/devdocs/reference/typing.html#numpy.typing.NDArray
beartype/beartype#96
https://github.com/patrick-kidger/torchtyping

cc @ezyang @malfet @rgommers @xuzhao9 @gramster

@mrshenli mrshenli added module: typing Related to mypy type annotations triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module labels Feb 25, 2022
@rgommers
Copy link
Collaborator

Thanks for this @JulesGM. The links contain very interesting discussion. My impression at the moment:

  • Adding dtype support ala numpy.typing.NDArray is feasible and a good idea at this time.
  • Shape support is much harder, and best left to third-party extensions at this time. I think PEP 646 actually skips over this (aside from a note that this is for the future: https://www.python.org/dev/peps/pep-0646/#shape-arithmetic); you have kinds similar needs as structured kernels and the meta backend in PyTorch - but then doing that with static analysis.
  • Named dimensions are useful and probably okay to add to PyTorch at this time. A little less sure about this one.

My strong belief is that the best thing to do now is to write a brand-new package for representing Array/Tensor types. Pretty much what the discussion in PEP 646 invites us to do, in fact.

This comment from beartype/beartype#96 sounds about right.

@JulesGM
Copy link
Author
JulesGM commented Mar 1, 2022

That would be really great, and I think it's exciting.

@ReHoss
Copy link
ReHoss commented May 15, 2025

Any news on this topic? Is jaxtyping the best solution for now?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
module: typing Related to mypy type annotations triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Development

No branches or pull requests

4 participants
0