code: def backward_(dz, x, mu, rs, gamma):
      wtype = gamma.dtype
      itype = x.dtype
      otype = dz.dtype
      ctype = mu.dtype
      mu = mu.unsqueeze(1)
      rs = rs.unsqueeze(1)
  
      hidden_size = gamma.numel()
      y = rs * (x.to(ctype) - mu)
      dbeta = dz.view(-1, hidden_size).sum(0, dtype=ctype)
      dgamma = (dz * y).view(-1, hidden_size).sum(0, dtype=ctype)
      dy = dz.view(-1, hidden_size).to(ctype) * gamma.unsqueeze(0).to(
          ctype
      )
      mdy = dy.mean(1, keepdim=True, dtype=ctype)
  
      mdyy = (dy * y).mean(1, keepdim=True, dtype=ctype)
      dx = rs * (dy - mdyy * y - mdy)
  
      return dx.to(itype), dgamma.to(wtype), dbeta.to(wtype)
  
file_path: apex/apex/contrib/test/layer_norm/test_fast_layer_norm.py
imports: import itertools
  import unittest
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: backward_(dz, x, mu, rs, gamma) and returns: dx,  dgamma,  dbeta