code: def partition_batch(input_) -> Tensor:
      tensor_parallel_mode = get_tensor_parallel_mode()
      if tensor_parallel_mode in _parallel_split_batch:
          if isinstance(input_, dict):
              return {
                  k: _parallel_split_batch[tensor_parallel_mode](v)
                  for k, v in input_.items()
              }
          else:
              return _parallel_split_batch[tensor_parallel_mode](input_)
      else:
          return input_
  
file_path: ColossalAI/colossalai/legacy/nn/layer/colossalai_layer/_utils.py
imports: import torch.nn
  from torch import Tensor
  from ..parallel_2d._operation import split_batch_2d
  from ..parallel_2p5d._operation import split_batch_2p5d
  from ..parallel_3d._operation import split_batch_3d
  from ..utils import get_tensor_parallel_mode
label_desc: Use image node_id 1 for calling a global function with example usage: partition_batch(input_) and returns: input_, unknown, _parallel_split_batch