PyTorch/LibTorch/TorchScript as a backend

So one thing I’m still thinking about is whether TorchScript functions is the right level or if I should try harder to only do individual ops.

Pro TorchScript:

  • Individual ops would go much deeper into the PyTorch JIT API than the very widely used TorchScript interface we use now.
  • We don’t have to figure out what to do with non-tensor arguments as they’ll be created as constants in TorchScript
  • We faithfully use the PyTorch chain, e.g. their optimizations for the graphs we feed.
  • There might be applications where having programs is actually better.

Pro Ops:

  • The match to Relay in granularity is much better and wrapping each op in a TorchScript function is quite an overhead.

What do you think?