Contributing to Tensor Shapes
Pyrefly's tensor shape tracking is designed so that coverage of the PyTorch library can be extended without understanding Pyrefly's internals. There are three ways to contribute:
- Fixture stubs —
.pyifiles with shape-generic type signatures for PyTorch modules and functions (e.g.,nn.Linear,torch.mm). - DSL functions — shape transform specifications for operators with
complex shape logic (e.g.,
reshape,cat,F.interpolate). - Ported models — fully annotated ports of real-world PyTorch models that serve as tests and examples.
For detailed instructions on writing stubs, DSL functions, and porting models, see the Tensor Shapes Contributing Guide in the repository. The stub library author reference documents the annotation constructs and DSL used by shape-aware stubs.