Fix the decomposition of native_layer_norm_backward operation. (#674)

The return type for native_layer_norm_backward is a tuple of
optional tensors as the output_mask decides whether or not the
gradient of a particular input is to be calculated or not.

Signed-Off-By: Prateek Gupta <[email protected]>

Author: gprateek93

functorch/_src/decompositions.py

@@ -441,7 +441,7 @@ def addmm(self: Tensor, mat1: Tensor, mat2: Tensor, beta: int = 1, alpha: int =
 
 
 @register_decomposition(aten.native_layer_norm_backward)
-def native_layer_norm_backward(grad_out: Tensor, input: Tensor, normalized_shape: List[int], mean: Tensor, rstd: Tensor, weight: Optional[Tensor], bias: Optional[Tensor], output_mask: List[bool]) -> Tuple[Tensor, Tensor, Tensor]:
+def native_layer_norm_backward(grad_out: Tensor, input: Tensor, normalized_shape: List[int], mean: Tensor, rstd: Tensor, weight: Optional[Tensor], bias: Optional[Tensor], output_mask: List[bool]) -> Tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
     input_shape = input.shape
     input_ndim = input.dim()
 

test/test_ops.py

@@ -1555,7 +1555,7 @@ def get_names(inpt):
                 f.write(f'{op}\n')
 
     def test_decompositions_torchscriptable(self, device):
-        skip_list = [torch.ops.aten.native_layer_norm_backward.default]
+        skip_list = []
         for op, decomposition in decomposition_table.items():
             if op in skip_list:
                 continue