code: def test_moe_gather(
      n_tokens: int, n_experts: int, n_top_k: int, do_padding: bool
  ):
      get_accelerator().manual_seed(0xC0FFEE)
  
      (
          batch,
          moe_input,
          scores,
          mapped_slots,
          expert_counts,
      ) = build_inputs(n_tokens, n_experts, n_top_k, do_padding)
  
      output = torch.randn(
          (batch.tensor_toks, 4096),
          dtype=torch.float16,
          device=get_accelerator().current_device(),
      )
  
      gather = MoEGather(DtypeEnum.fp16, 4096)
      gather(output, moe_input, scores, mapped_slots, expert_counts)
  
      for token_idx in range(n_tokens):
          effective_score = scores[token_idx].sum().item()
          assert torch.equal(
              output[token_idx],
              torch.full(
                  (4096,),
                  token_idx * effective_score,
                  dtype=torch.float16,
                  device=get_accelerator().current_device(),
              ),
          )
  
file_path: DeepSpeed/tests/unit/inference/v2/kernels/ragged_ops/test_moe_gather.py
imports: import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.inference_utils import DtypeEnum
  from deepspeed.inference.v2.kernels.ragged_ops import MoEGather, MoEScatter, RaggedTopKGating
  from .ragged_testing_utils import build_simple_batch
label_desc: Use image node_id 2 for calling a global function with example usage: test_moe_gather(n_tokens, n_experts, n_top_k, do_padding) without return types