code: def forward(ctx, x, gamma, beta, epsilon, memory_efficient=False):
      ctx.x_shape = x.shape
      ctx.memory_efficient = memory_efficient
  
      x = x.contiguous()
      gamma = gamma.contiguous()
      beta = beta.contiguous()
      hidden_size = gamma.numel()
      xmat = x.view((-1, hidden_size))
      ymat, mu, rsigma = fast_layer_norm.ln_fwd(
          xmat, gamma, beta, epsilon
      )
      if ctx.memory_efficient:
          ctx.save_for_backward(ymat, gamma, None, rsigma, beta)
      else:
          ctx.save_for_backward(xmat, gamma, mu, rsigma, None)
      return ymat.view(x.shape)
  
file_path: apex/apex/contrib/layer_norm/layer_norm.py
imports: import torch
  from torch.nn import init
  from apex._autocast_utils import _cast_if_autocast_enabled
  import fast_layer_norm
label_desc: Use image node_id 1 for calling the FastLayerNormFN obj's underlying member method code with example usage: obj.forward(ctx, x, gamma, beta, epsilon, memory_efficient) and returns: ymat