patrick-kidger / jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/
Other
1.05k stars 49 forks source link

Are pytorch named tensors supported, like in torchtyping? #200

Open zxti opened 3 months ago

patrick-kidger commented 3 months ago

I'm afraid there is no special support for these. I don't really recommend using them -- PyTorch itself seems to have mostly given up on adding further support for them. So they're only supported on a small number of PyTorch operations anyway.