code: def all_gather_tensor_3d(
      tensor: Tensor, dim: int, parallel_mode: ParallelMode
  ) -> Tensor:
      r"""All-reduce the gradient in backward pass.
  
      Args:
          tensor (:class:`torch.tensor`): Input tensor.
          dim (int): Dimension to gather.
          parallel_mode (:class:`colossalai.legacy.context.parallel_mode.ParallelMode`): 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 _AllGatherTensor3D.apply(tensor, dim, 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 10 for calling a global function with example usage: all_gather_tensor_3d(tensor, dim, parallel_mode) and returns: _AllGatherTensor3D