so recently, I looked a bit into TorchScript usability and tried to collect some ideas in the ScriptTorch blog post.
I see two main things that I would like to see improved:
User experience: Easier scripting of models (e.g. by having fallbacks).
- I implemented some bits on a github branch and will push a detailed description of a possible fallback implementation. This covers functions/datatypes but not syntax.
- I have some ideas how one could implement lambdas (by inlining in the ConvertToSSA step and re-doing type checks or just re-compile to cope with the binding behaviour) used within torchscript.
I want to enable doing graph transformations from Python. I’m toying around with it in a blog post (note that this isn’t intended to be a serious vmap implementation, that is just an example).
- Safety (no segfaults): Graceful invalidation of Python Node/Value/Block when C++ object is deleted by t-vi · Pull Request #50326 · pytorch/pytorch · GitHub
- API exposure. I have some patches on the branch linked above, too, that I’d upstream bit by bit.
A first cleanup is Merge CompilationUnit from torch._C and torch.jit by t-vi · Pull Request #50614 · pytorch/pytorch · GitHub but there is more.
- Utility functions. Probably it would make sense to mark these experimental for a while.
I already got quite a bit of feedback on slack and in-person calls (thank you!), and I’ll try to summarize things a bit, but maybe here is a good place for additional feedback.
Given that this is quite a large program, it would be cool to get your input on both what’s missing and what how you would tackle them (in terms of prioritization and also implementation).