code: def drop_tokens(input_, dim=0):
      mpu = deepspeed.utils.groups.mpu
      if mpu is None or mpu.get_tensor_model_parallel_world_size() == 1:
          # no tensor parallelism for non-experts
          return input_
      return _DropTokens.apply(input_, dim)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/moe/mappings.py
imports: import torch
  import deepspeed
label_desc: Use image node_id 4 for calling a global function with example usage: drop_tokens(input_, dim) and returns: _DropTokens, input_