Add decomposition for aten.native_layer_norm_backward op.

Signed-Off-By: Prateek Gupta <prateek@nod-labs.com>

Author: gprateek93

functorch/_src/decompositions.py

@@ -1,5 +1,5 @@
 import torch
-from torch import Tensor
+from torch import Tensor, reciprocal_
 from typing import Optional, List, Tuple
 from enum import Enum
 
@@ -381,6 +381,43 @@ def native_layer_norm(input: Tensor, normalized_shape: List[int], weight: Option
     return (out, mean, rstd)
 
 
+@register_decomposition(aten.native_layer_norm)
+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]:
+    input_shape = input.shape
+    input_ndim = input.dim()
+
+    axis = input_ndim - len(normalized_shape)
+    inner_dims = input_shape[axis:]
+    outer_dims = input_shape[:axis]
+    N = prod(inner_dims)
+
+    x_hat = (input - mean) * rstd
+    grad_x_hat = []
+    if weight is not None:
+        grad_x_hat = grad_out * weight
+    else:
+        grad_x_hat = grad_out
+    a = grad_x_hat * N
+    b = aten.sum(grad_x_hat, inner_dims, True)
+    c1 = grad_x_hat * x_hat
+    c2 = aten.sum(c1, inner_dims, True)
+    c3 = x_hat * c2
+
+    inner = (a - b) - c3
+
+    d_weight = []
+    d_bias = []
+    d_input = []
+
+    if output_mask[0]:
+        d_input = (rstd / N ) * inner
+    if output_mask[1]:
+        d_weight = aten.sum((grad_out * x_hat), outer_dims, False)
+    if output_mask[2]:
+        d_bias = aten.sum(grad_out, outer_dims, False)
+
+    return (d_input, d_weight, d_bias)
+
 # @register_decomposition(aten.addmm)
 # def addmm(self: Tensor, mat1: Tensor, mat2: Tensor, beta=1, alpha=1):
 #     if not self.is_floating_point():