code: def _gather_tokens(input_, dim=0):
      """Gather tensors and concatenate them along a dimension"""
      mpu = deepspeed.utils.groups.mpu
  
      input_ = input_.contiguous()
      # Size and dimension.
      rank = mpu.get_tensor_model_parallel_rank()
  
      tensor_list = [
          torch.empty_like(input_)
          for _ in range(mpu.get_tensor_model_parallel_world_size())
      ]
      tensor_list[rank] = input_
      deepspeed.comm.all_gather(
          tensor_list,
          input_,
          group=mpu.get_tensor_model_parallel_group(),
      )
  
      # Note: torch.cat already creates a contiguous tensor.
      output = torch.cat(tensor_list, dim=dim).contiguous()
  
      return output
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/moe/mappings.py
imports: import torch
  import deepspeed
label_desc: Use image node_id 1 for calling a global function with example usage: _gather_tokens(input_, dim) and returns: output