code: def _split(input_, dim=-1, process_group=None):
      # skip if only one rank involved
      world_size = dist.get_world_size(process_group)
      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 = dist.get_rank(process_group)
      output = tensor_list[rank].clone().contiguous()
  
      return output
  
imports: import torch
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 3 for calling a global function with example usage: _split(input_, dim, process_group) and returns: output, input_