code: def layernorm_2d(
      input_: Tensor,
      E_x: Tensor,
      Var_x: Tensor,
      hidden_size: int,
      row_parallel_mode: ParallelMode,
      col_parallel_mode: ParallelMode,
  ) -> Tensor:
      r"""Layernorm.
  
      Args:
          input_ (:class:`torch.tensor`): input matrix.
          E_x (:class:`torch.tensor`): mean.
          Var_x (:class:`torch.tensor`): variance.
          hidden_size (int): hidden size.
          row_parallel_mode (:class:`colossalai.legacy.context.ParallelMode`): row parallel mode.
          col_parallel_mode (:class:`colossalai.legacy.context.ParallelMode`): column parallel mode.
  
      Note:
          The parallel_mode should be concluded in ``ParallelMode``. More details about ``ParallelMode`` could be found
          in `parallel_mode <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/context/parallel_mode.py>`_
      """
      return _Layernorm_2D.apply(
          input_,
          E_x,
          Var_x,
          hidden_size,
          row_parallel_mode,
          col_parallel_mode,
      )
  
imports: from typing import Any, Optional, Tuple
  import torch
  import torch.distributed
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication.collective import all_gather, all_reduce, reduce_scatter
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.utils import get_current_device
label_desc: Use image node_id 4 for calling a global function with example usage: layernorm_2d(input_, E_x, Var_x, hidden_size, row_parallel_mode, col_parallel_mode) and returns: _Layernorm_2D