code: def _drop_tokens(input_, dim=0):
      """Divide a tensor among the tensor parallel ranks"""
      mpu = deepspeed.utils.groups.mpu
  
      total_chunks = mpu.get_tensor_model_parallel_world_size()
      this_chunk = mpu.get_tensor_model_parallel_rank()
      assert (
          input_.shape[dim] % total_chunks == 0
      ), f"input dimension {dim} ({input_.shape[dim]}) is not divisible by tensor parallel world size ({total_chunks})"
      chunk_size = input_.shape[dim] // total_chunks
  
      return torch.narrow(
          input_, dim, this_chunk * chunk_size, chunk_size
      )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/moe/mappings.py
imports: import torch
  import deepspeed
label_desc: Use image node_id 2 for calling a global function with example usage: _drop_tokens(input_, dim) and returns: torch