TorchScript attempts to compile dead branch of torch.jit.is_scripting

[sharvil]

🐛 Describe the bug

Consider the following code. The else branch of the condition in get_context evaluates to False when scripting. Running this code still produces an exception.

import torch

meow = 1239

def get_context():
  if torch.jit.is_scripting():
    print('hi')
  else:
    global meow
    meow = 9123
    return meow

if __name__ == '__main__':
  x = torch.jit.script(get_context)
  print(x.code)

Output:

torch.jit.frontend.UnsupportedNodeError: global variables aren't supported:
  File "/mnt/data/lmnt/code/tacotron/src/lmnt/tacotron/context.py", line 9
    print('hi')
  else:
    global meow
    ~~~~~~ <--- HERE
    meow = 9123
    return meow

This is the output of the script after removing the global meow line:

def get_context() -> NoneType:
  print("hi")
  return None

Note how the dead branch is eliminated after scripting. Instead, the dead branch should be eliminated before attempting to compile it.

There are workarounds to the current behavior (e.g. move the else clause into a separate function with a @torch.jit.unused annotation) but they aren't ergonomic and often ruin the code flow.

Versions

Any version of PyTorch

[Gamrix]

@eellison This seems very useful, but also hard to do as the error gets thrown while we are still building the AST. Is there anything we can do here, like special casing torch.jit.is_scripting() and inlining if statements ahead of time?

[sharvil]

I like the idea of special-casing torch.jit.is_scripting() until a more general solution is available.

Perhaps a more general way to handle this issue is to support constant folding and constant propagation and evaluate constant conditionals during the AST building pass.