code: def test_determinism():
      """
      Ensure that ties between two logits are broken deterministically. This is essential when
      the gating is distributed across multiple devices that need to map the same token to
      the same expert.
      """
  
      n_tokens = 512
      n_experts = 64
      n_top_k = 1
  
      logits = torch.zeros(
          (n_tokens, n_experts),
          dtype=torch.float16,
          device=get_accelerator().current_device(),
      )
      batch = build_simple_batch([n_tokens], padding=False)
  
      logits[:, 19] = 1.0
      logits[:, 26] = 1.0
  
      gate = RaggedTopKGating(DtypeEnum.fp16)
  
      for _ in range(1024):
          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[19] == n_tokens
          assert expert_counts[26] == 0
          assert torch.all(expert_assignment == 19)
          assert torch.unique(expert_offset).shape[0] == n_tokens
          assert allclose(
              scores,
              F.softmax(logits.float(), dim=1)[:, 19].reshape(-1, 1),
          )
  
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 6 for calling a global function with example usage: test_determinism() without return types