functorch icon indicating copy to clipboard operation
functorch copied to clipboard

backwards decomposition returns None, torchscript doesn't convert to undefined tensor -> Error from autograd

Open samdow opened this issue 3 years ago • 0 comments

The decomposition of layer_norm_backward returns None. When AOTAutograd encounters this, torch_dispatch will convert that to an undefined Tensor, getting correct behavior. However, TorchScript will still return a None and doesn't have a great way to distinguish this function from a torchscripted function that should actually be returning None. Then, autograd expects everything on the stack to be a Tensor and errors when it sees the None instead

Short term paths forward: (1) decomposition in functorch returns zeros (this will probably go in so that we have it and can test against it. Happy to switch to another path after if people have thoughts) (2) port this decomposition to C++ and use undefined tensor there (3?) Figure out if there's a way to ask for an undefined tensor in Python that Torchscript will see and understand

Long term paths forward: (1) Torchscript converts Nones to undefined tensors in specific cases where (maybe) it's a decomposition of a backwards function (2) Stick with the original short term path forward?

cc @Chillee @eellison

samdow avatar May 17 '22 21:05 samdow