code: def layernorm_3d(
      input_: Tensor,
      weight: Tensor,
      bias: Tensor,
      normalized_shape: int,
      eps: float,
      output_parallel_mode: ParallelMode,
      input_x_weight_parallel_mode: ParallelMode,
  ) -> Tensor:
      r"""3D parallel Layernorm.
  
      Args:
          input_ (:class:`torch.tensor`): input matrix.
          weight (:class:`torch.tensor`): matrix of weight.
          bias (:class:`torch.tensor`): matrix of bias.
          normalized_shape (int): input shape from an expected input of size.
              :math:`[* \times \text{normalized_shape}[0] \times \text{normalized_shape}[1]
              \times \ldots \times \text{normalized_shape}[-1]]`
              If a single integer is used, it is treated as a singleton list, and this module will
              normalize over the last dimension which is expected to be of that specific size.
          eps (float): a value added to the denominator for numerical stability
          output_parallel_mode (:class:`colossalai.legacy.context.parallel_mode.ParallelMode`): output parallel mode.
          input_x_weight_parallel_mode (:class:`colossalai.legacy.context.parallel_mode.ParallelMode`): input x weight 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 _Layernorm3D.apply(
          input_,
          weight,
          bias,
          id(weight),
          id(bias),
          normalized_shape,
          eps,
          output_parallel_mode,
          input_x_weight_parallel_mode,
      )
  
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication import all_gather, all_reduce, broadcast, reduce, reduce_scatter
  from colossalai.legacy.constants import INPUT_GROUP_3D, WEIGHT_GROUP_3D
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from ._utils import get_parallel_mode_from_env, push_async_grad
label_desc: Use image node_id 6 for calling a global function with example usage: layernorm_3d(input_, weight, bias, normalized_shape, eps, output_parallel_mode, input_x_weight_parallel_mode) and returns: _Layernorm3D