code: def build_inputs(
      n_tokens: int, n_experts: int, n_top_k: int, do_padding: bool
  ):
      assert n_tokens <= 2048, "This test will break if n_tokens > 2048"
  
      # Sequence composition shouldn't matter here
      batch = build_simple_batch([n_tokens], padding=do_padding)
  
      logits = torch.randn(
          (batch.tensor_toks, n_experts),
          dtype=torch.float16,
          device=get_accelerator().current_device(),
      )
  
      # This will make each token's value equal to its index. NOTE: This will break for
      # tokens with index > 2048.
      hidden_states = (
          torch.arange(
              batch.tensor_toks,
              dtype=torch.float16,
              device=get_accelerator().current_device(),
          )
          .repeat_interleave(4096, dim=0)
          .reshape(batch.tensor_toks, 4096)
          .contiguous()
      )
  
      gate = RaggedTopKGating(DtypeEnum.fp16)
  
      # Gating outputs
      expert_counts = torch.zeros(
          (n_experts,),
          dtype=torch.int32,
          device=get_accelerator().current_device(),
      )
      scores = torch.empty(
          (batch.tensor_toks, n_top_k),
          dtype=torch.float32,
          device=get_accelerator().current_device(),
      )
      expert_assignment = torch.empty(
          (batch.tensor_toks, n_top_k),
          dtype=torch.int32,
          device=get_accelerator().current_device(),
      )
      expert_offset = torch.empty(
          (batch.tensor_toks, n_top_k),
          dtype=torch.int32,
          device=get_accelerator().current_device(),
      )
  
      gate(
          expert_counts,
          scores,
          expert_assignment,
          expert_offset,
          logits,
          batch,
      )
  
      # Scatter outputs
      moe_input = torch.empty(
          (batch.tensor_toks * n_top_k, 4096),
          dtype=torch.float16,
          device=get_accelerator().current_device(),
      )
      expert_cumsum = torch.empty(
          (n_experts,),
          dtype=torch.int64,
          device=get_accelerator().current_device(),
      )
      mapped_slots = torch.empty(
          (batch.tensor_toks, n_top_k),
          dtype=torch.int32,
          device=get_accelerator().current_device(),
      )
  
      scatter = MoEScatter(DtypeEnum.fp16, 4096)
      scatter(
          moe_input,
          expert_cumsum,
          mapped_slots,
          hidden_states,
          expert_counts,
          expert_assignment,
          expert_offset,
      )
  
      return batch, moe_input, scores, mapped_slots, expert_counts
  
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 1 for calling a global function with example usage: build_inputs(n_tokens, n_experts, n_top_k, do_padding) and returns: batch,  moe_input,  scores,  mapped_slots,  expert_counts