TorchScript compatibility

Open fferroni opened this issue 2 years ago • 1 comments

Hi,

The package now does not require JAX but can also work with torch.Tensor I was trying this, and it worked fine, but it fails when using torch.jit.script In TorchTyping you fixed it with https://github.com/patrick-kidger/torchtyping/issues/13 Is there something similar here?

Feb 09 '23 19:02 fferroni

Interesting point! Supportig this is definitely a non-goal for jaxtyping, although I'd be happy to accept PRs if something needs tweaking to enable torchscript compatibility.

For what it's worth, my impression is that TorchScript is going to be deprecated at some point in favour of PyTorch 2's new compiler offerings (TorchDynamo etc.)

Feb 13 '23 21:02 patrick-kidger