Tensor shapes beyond PyTorch

[EtaoinWu]

Hello, the awesome team at Meta! Currently the tensor shapes feature is only available and documented for PyTorch. Is there any plans to support it for other array libraries, like NumPy or JAX? I'm a heavy user of JAX and jaxtyping mostly for scientific competition work (non ML), and really hope to see the support on there. Thank you!

[EtaoinWu]

See also #2332

[rchen152]

Hey @EtaoinWu - we do support jaxtyping annotations (https://pyrefly.org/en/docs/tensor-shapes/#jaxtyping-runtime-type-checking), with the caveat that tensor shape checking is experimental and changing quickly, including the jaxtyping frontend. I believe @stroxler is also looking to introduce numpy compatibility of some sort, so stay tuned :)

[jorenham]

like NumPy

NumPy has preliminary support for type-checker agnostic shape-typing out-of-the box, and this has improved a lot in the latest v2.5.0 release.

But due to limitations in Python's type system, it's not possible to meaningfully type the individual axis sizes, only the type of the shape as a whole (i.e. as a tuple of plain ints), so for those in particular, Pyrefly's tensor types might be useful in certain cases.