Torch.jit.is_scripting and handle_torch_function

jit
1

Hello,

I observed in PyTorch’s source code some functions that follow the pattern below (for example this one):

def my_function(a, b, c, d=None, e=None):
    if not torch.jit.is_scripting():
        tens_ops = (a, b, c)
        if any([type(t) is not Tensor for t in tens_ops]) and has_torch_function(tens_ops):
            return handle_torch_function(my_function , a, b, c, d=d, e=e)
   # something else here

Could anybody please tell me what’s going on there? In particular, what do torch.jit.is_scripting and handle_torch_function do?

And more importantly: Where can I learn those things in detail so that I can myself write similar custom functions and stop asking stupid questions like these? :stuck_out_tongue: It seems that these are about TorchScript, but I wanted to ask just in case, before diving into that.

Thank you very much in advance for your help!

2

torch.jit.is_scripting() can be used as a guard for Python-only methods, which are not exportable.
Since the JIT cannot export all arbitrary Python objects, you could thus use different “paths” in your code for the Python execution and the scripting of the mdoel.
handle_torch_function sounds like a custom function definition.

As these features are currently being developed, the documentation and tutorials might not be up to date, so please ask these kind of questions here. :wink:

torch.jit.is_scripting can be found in the current master docs, but might also be updated before the next release.

2 Likes
3

Thanks a lot, @ptrblck! Your answer is very helpful, as always.

1 Like