TorchScript usability

[tom]

Hi,

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.

• Dev experience:
  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).

Best regards

Thomas

[tom]

It’s too late to edit, but the broken link should be “I’m toying around with it in a blog post”. Thank you, Yanan Cao, for pointing it out!

[tom]

As a follow-up on the general usability topic: I wrote a new blog post about a potential route to implement fallbacks to Python: https://lernapparat.de/jit-fallback/
This only covers datatypes and functions the JIT cannot handle, not syntax constructs, but I think it might be a good way to give the user more freedom to use whatever tools they see fit (but obviously, the JIT won’t know what’s happening in them).