code: def _split(input_, parallel_mode, dim=-1):
      # skip if only one rank involved
      world_size = gpc.get_world_size(parallel_mode)
      if world_size == 1:
          return input_
  
      # Split along last dimension.
      dim_size = input_.size(dim)
      assert dim_size % world_size == 0, (
          f"The dimension to split ({dim_size}) is not a multiple of world size ({world_size}), "
          f"cannot split tensor evenly"
      )
  
      tensor_list = torch.split(input_, dim_size // world_size, dim=dim)
      rank = gpc.get_local_rank(parallel_mode)
      output = tensor_list[rank].contiguous()
  
      return output
  
imports: import torch
  import torch.distributed
  from colossalai.legacy.core import global_context
  from colossalai.legacy.global_variables import tensor_parallel_env
  from ..utils import divide
label_desc: Use image node_id 6 for calling a global function with example usage: _split(input_, parallel_mode, dim) and returns: output, input_