class_docstr: CUDA implementation of MoE gather. This will bring the tokens back
  to their original indices and perform the output scaling.
docstr: Reorders the moe_output tokens into their original order and scales them by their
  gating scale. This will be a no-op for padded tokens.
  
  Arguments:
      layer_output (torch.Tensor): The output of the layer of shape [n_tokens, hidden_size]. This has been scaled appropriately.
      moe_output (torch.Tensor): The output of the MoE of shape [n_tokens * n_top_k, hidden_size].
      scores (torch.Tensor): The gating scores of shape [n_tokens].
      mapped_slots (torch.Tensor): The index of the token in the expert's input of shape [n_tokens, n_top_k]. The indices of token ``i`` in layer_output is ``mapped_slots[i]``.
      expert_counts (torch.Tensor): The number of tokens assigned to each expert of shape [n_experts]. This is passed to fuse the clearing of this data structure into the gather.
  
  Returns:
      layer_output
docstr_tok: ["Reorders","the","moe_output","tokens","into","their","original","order","and","scales","them","by","their","gating","scale",".","This","will","be","a","no-op","for","padded","tokens",".","Arguments",":","layer_output","(","torch.Tensor",")",":","The","output","of","the","layer","of","shape","[","n_tokens",",","hidden_size","]",".","This","has","been","scaled","appropriately",".","moe_output","(","torch.Tensor",")",":","The","output","of","the","MoE","of","shape","[","n_tokens","*","n_top_k",",","hidden_size","]",".","scores","(","torch.Tensor",")",":","The","gating","scores","of","shape","[","n_tokens","]",".","mapped_slots","(","torch.Tensor",")",":","The","index","of","the","token","in","the","expert","'s","input","of","shape","[","n_tokens",",","n_top_k","]",".","The","indices","of","token","``","i","``","in","layer_output","is","``","mapped_slots","[","i","]","``",".","expert_counts","(","torch.Tensor",")",":","The","number","of","tokens","assigned","to","each","expert","of","shape","[","n_experts","]",".","This","is","passed","to","fuse","the","clearing","of","this","data","structure","into","the","gather",".","Returns",":","layer_output"]
code: def __call__(
      self,
      layer_output: torch.Tensor,
      moe_output: torch.Tensor,
      scores: torch.Tensor,
      mapped_slots: torch.Tensor,
      expert_counts: torch.Tensor,
  ) -> torch.Tensor:
      """
      Reorders the moe_output tokens into their original order and scales them by their
      gating scale. This will be a no-op for padded tokens.
  
      Arguments:
          layer_output (torch.Tensor): The output of the layer of shape [n_tokens, hidden_size]. This has been scaled appropriately.
          moe_output (torch.Tensor): The output of the MoE of shape [n_tokens * n_top_k, hidden_size].
          scores (torch.Tensor): The gating scores of shape [n_tokens].
          mapped_slots (torch.Tensor): The index of the token in the expert's input of shape [n_tokens, n_top_k]. The indices of token ``i`` in layer_output is ``mapped_slots[i]``.
          expert_counts (torch.Tensor): The number of tokens assigned to each expert of shape [n_experts]. This is passed to fuse the clearing of this data structure into the gather.
  
      Returns:
          layer_output
      """
      self.kernel(
          layer_output,
          moe_output,
          scores,
          mapped_slots,
          expert_counts,
          self.normalize_scores,
      )
      return layer_output
  
file_path: deepspeed/deepspeed/inference/v2/kernels/ragged_ops/moe_gather/moe_gather.py
imports: import torch
  from ...None import DSKernelBase
  from ....inference_utils import DtypeEnum
  from deepspeed.ops.op_builder import RaggedOpsBuilder
label_desc: Use image node_id 2 for calling the MoEGather obj's underlying member method code with example usage: obj.__call__(layer_output, moe_output, scores, mapped_slots, expert_counts) and returns: layer_output