code: def _top_k_gating_testing_helper(
      n_tokens: int, n_experts: int, n_top_k: int, seed: int = 0xC0FFEE
  ) -> None:
      torch.manual_seed(seed)
      logits = torch.randn(
          (n_tokens, n_experts),
          dtype=torch.float16,
          device=get_accelerator().current_device(),
      )
      batch = build_simple_batch([n_tokens], padding=False)
      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(),
      )
  
      gate(
          expert_counts,
          scores,
          expert_assignment,
          expert_offset,
          logits,
          batch,
      )
  
      ref_weights = F.softmax(logits, dim=-1, dtype=torch.float32)
      ref_scores, ref_indices = torch.topk(ref_weights, n_top_k, dim=-1)
  
      assert allclose(
          scores, ref_scores
      ), f"expected {ref_scores}, got {scores}"
      assert torch.equal(
          expert_assignment, ref_indices.to(torch.int32)
      ), f"expected {ref_indices}, got {expert_assignment}"
      assert (
          expert_counts.sum() == n_tokens * n_top_k
      ), f"expected {n_tokens * n_top_k} tokens, got {expert_counts.sum()}"
  
      # Ensure that the expert offsets are unique
      for i in range(n_experts):
          expert_idxs = torch.where(
              expert_assignment == i, expert_offset, 0
          )
          if expert_counts[i] > 0:
              assert (
                  expert_idxs.unique().shape[0] == expert_counts[i]
              ), f"expected {expert_counts[i]} unique offsets, got {expert_idxs.unique().shape[0]}"
              assert (
                  expert_idxs.max() == expert_counts[i] - 1
              ), f"expected max offset {expert_counts[i] - 1}, got {expert_idxs.max()}"
          else:
              # Should have all 0's so one unique value
              assert expert_idxs.unique().shape[0] == 1
              assert expert_idxs.max() == 0
  
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 1 for calling a global function with example usage: _top_k_gating_testing_helper(n_tokens, n_experts, n_top_k, seed) without return types