I’ve just opened a Github Issue relating to type checking in Pytorch. I wanted to start a discussion around the topic here as well.
Here is the content of the issue:
Currently, mypy is used to typecheck PyTorch, with lint runner and dmypy. We appreciate the community’s work maintaining mypy and type coverage in PyTorch and want to build on that foundation. Pyrefly is a new standards-compliant Python type checker. The Pyrefly team has been hard at work on building a performant and backwards compatible type checker, that we think can improve the current type checking setup and the development experience for PyTorch users.
For example, the current setup can make it tricky to know which files are being typechecked and in which mode (strict vs. default). Use of type: ignore and # mypy: ignore-errors to disable checks on certain files, adds to the challenge. We think we can make typechecking simpler, and improve the overall experience.
We are proposing that we add Pyrefly as a typechecker to Pytorch.
The benefits to PyTorch will be:
Whole repo checks in seconds:
Eliminates inconsistencies created by dmypy between local and CI revisions
Fast CI signal
A richer IDE experience with types that match hover and diagnostics. We support most major IDEs: IDE Installation | Pyrefly
Modern type checking features:
Container inference, conformance to the typing spec for feature support
We’ve found over 200 bugs in PyTorch just by enabling Pyrefly for testing!
Clear configuration and ownership:
The Pyrefly team will help support typing in PyTorch and respond quickly to issues on our github
Here’s how we would propose to get Pyrefly up and running in PyTorch:
Phase 1:
Check in Pyrefly configs along with the suppressions needed for Pyrefly to check cleanly
Add a non-blocking CI linter runner job to observe changes and test the integration
Gather community feedback on the checker and IDE extension via:
Download the VSCode Extension
Run pyrefly through lint runner
Phase 2:
With community buy in, we’ll swap the checker from mypy to pyrefly over the weekend to be least disruptive
After Pyrefly has been enabled smoothly for a few days, we’ll cleanup the unused type: ignores that remain in the code. We plan to cleanup approximately 600 unused mypy ignores
Phase 3:
Work with the community to help add types where they are useful in PyTorch and answer questions around typing features and usage
Set up additional jobs like pyrefly infer to ease the process of adding types
Work with the community to better export types to consumers of PyTorch
We’d love for you to try Pyrefly, share your experiences, and help us make it (and PyTorch) even better
