code: def _test_single_mapping_helper(
      n_tokens: int,
      n_experts: int,
      assigned_expert: int,
      logit_fill: float = 0.0,
      match_fill: float = 1.0,
  ) -> None:
      n_top_k = 1
      logits = torch.full(
          (n_tokens, n_experts),
          logit_fill,
          dtype=torch.float16,
          device=get_accelerator().current_device(),
      )
  
      logits[:, assigned_expert] = match_fill
  
      gate = RaggedTopKGating(DtypeEnum.fp16)
  
      expert_counts = torch.zeros(
          (n_experts,),
          dtype=torch.int32,
          device=get_accelerator().current_device(),
      )
      scores = torch.empty(
          (n_tokens, n_top_k),
          dtype=torch.float32,
          device=get_accelerator().current_device(),
      )
      expert_assignment = torch.empty(
          (n_tokens, n_top_k),
          dtype=torch.int32,
          device=get_accelerator().current_device(),
      )
      expert_offset = torch.empty(
          (n_tokens, n_top_k),
          dtype=torch.int32,
          device=get_accelerator().current_device(),
      )
      batch = build_simple_batch([n_tokens], padding=False)
  
      gate(
          expert_counts,
          scores,
          expert_assignment,
          expert_offset,
          logits,
          batch,
      )
  
      assert expert_counts[assigned_expert] == n_tokens
      assert torch.all(expert_assignment == assigned_expert)
      assert torch.unique(expert_offset).shape[0] == n_tokens
      assert allclose(
          scores,
          F.softmax(logits.float(), dim=1)[:, assigned_expert].reshape(
              -1, n_top_k
          ),
      )
  
file_path: DeepSpeed/tests/unit/inference/v2/kernels/ragged_ops/test_top_k_gating.py
imports: import pytest
  import torch
  import torch.nn.functional
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.inference_utils import DtypeEnum
  from deepspeed.inference.v2.kernels.ragged_ops import RaggedTopKGating
  from .ragged_testing_utils import build_simple_batch
  from ...inference_test_utils import allclose
label_desc: Use image node_id 3 for calling a global function with example usage: _test_single_mapping_helper(n_tokens, n_experts, assigned_expert, logit_fill, match_fill) without return types