code: def forward(ctx, inputs, normalized_shape, weight, bias, eps=1e-05):
      if (
          not inputs.is_contiguous()
          and normalized_shape != inputs.shape[-1:]
      ):
          raise ValueError(
              f"This implementation only support normalizing along the last dimension for "
              f"noncontiguous inputs. I.e., we expect "
              f"normalized_shape={tuple(inputs.shape[-1:])}, but got {normalized_shape} instead"
          )
      if not inputs.is_contiguous() and inputs.dim() != 4:
          raise ValueError(
              f"This implementation only supports 4-dim noncontiguous inputs, but got "
              f"{inputs.dim()} instead"
          )
  
      normalized_degree = len(normalized_shape)
      layer_shape = inputs.shape[:-normalized_degree]
      M, N = prod(layer_shape), prod(normalized_shape)
  
      x_invstd = torch.empty(
          M, dtype=torch.float32, device=inputs.device
      )
      x_mean = torch.empty(M, dtype=torch.float32, device=inputs.device)
      y = torch.empty(
          inputs.shape, dtype=inputs.dtype, device=inputs.device
      )
  
      grid = lambda kwargs: (
          triton.cdiv(kwargs["M"], kwargs["M_BLOCK"]),
      )
      if inputs.is_contiguous():
          _layer_norm_forward[grid](
              x_ptr=inputs,
              w_ptr=weight,
              b_ptr=bias,
              eps=eps,
              x_invstd_ptr=x_invstd,
              x_mean_ptr=x_mean,
              y_ptr=y,
              M=M,
              N=N,
          )
      else:
          D0, D1, D2, D3 = inputs.shape
          S0, S1, S2, S3 = inputs.stride()
          _layer_norm_forward_strided[grid](
              x_ptr=inputs,
              w_ptr=weight,
              b_ptr=bias,
              eps=eps,
              x_invstd_ptr=x_invstd,
              x_mean_ptr=x_mean,
              y_ptr=y,
              M=M,
              N=N,
              D0=D0,
              D1=D1,
              D2=D2,
              D3=D3,
              S0=S0,
              S1=S1,
              S2=S2,
              S3=S3,
          )
  
      ctx.save_for_backward(inputs, weight, x_invstd, x_mean)
      ctx.flatten_shape = M, N
      return y
  
file_path: apex/apex/contrib/openfold_triton/layer_norm.py
imports: from math import prod
  import torch
  import triton
  from torch.autograd import Function
  from apex.contrib.openfold_triton._layer_norm_backward_kernels import PARTIAL_REDUCE_MIN, _layer_norm_backward_buf_reduce, _layer_norm_backward_dw_db_partial, _layer_norm_backward_dw_db_partial_strided, _layer_norm_backward_dx, _layer_norm_backward_dx_strided
  from apex.contrib.openfold_triton._layer_norm_forward_kernels import _layer_norm_forward, _layer_norm_forward_strided
label_desc: Use image node_id 1 for calling the LayerNormSmallShapeOptImpl obj's underlying member method code with example usage: obj.forward(ctx, inputs, normalized_shape, weight, bias, eps) and returns: y