code: def split_batch_3d(
      input_: Tensor,
      dim: int = 0,
      input_parallel_mode: ParallelMode = ParallelMode.PARALLEL_3D_INPUT,
      weight_parallel_mode: ParallelMode = ParallelMode.PARALLEL_3D_WEIGHT,
  ) -> Tensor:
      r"""Splits 3D tensor in batch.
  
      Args:
          input_ (:class:`torch.tensor`): Input tensor.
          dim (int): Specified dimension in which to split.
          input_parallel_mode (:class:`colossalai.legacy.context.parallel_mode.ParallelMode`, optional): input parallel mode.
          weight_parallel_mode (:class:`colossalai.legacy.context.parallel_mode.ParallelMode`, optional): weight parallel mode.
  
      Returns:
          :class:`torch.tensor`: The tensor has been split.
  
      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>`_.
      """
      if input_.size(dim) <= 1:
          return input_
      weight_parallel_mode = get_parallel_mode_from_env(WEIGHT_GROUP_3D)
      input_parallel_mode = get_parallel_mode_from_env(INPUT_GROUP_3D)
      weight_world_size = gpc.get_world_size(weight_parallel_mode)
      input_world_size = gpc.get_world_size(input_parallel_mode)
      output = torch.chunk(input_, weight_world_size, dim=dim)[
          gpc.get_local_rank(weight_parallel_mode)
      ].contiguous()
      output = torch.chunk(output, input_world_size, dim=dim)[
          gpc.get_local_rank(input_parallel_mode)
      ].contiguous()
      return output
  
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 8 for calling a global function with example usage: split_batch_3d(input_, dim, input_parallel_mode, weight_parallel_mode) and returns: output, input_