code: def test_problem_size_permutations(model_dim: int) -> None:
      """
      Validate for different embedding sizes.
      """
      dtype = torch.float16
      seq_lens = [128, 64, 192, 32]
  
      batch = build_simple_batch(seq_lens, padding=False)
      hidden_states = torch.randn(
          (batch.current_tokens, model_dim),
          dtype=dtype,
          device=get_accelerator().current_device(),
      )
  
      reference_result = baseline_implementation(
          hidden_states, seq_lens
      )
  
      kernel = RaggedLogitsGather(model_dim, dtype)
      output = torch.empty_like(reference_result)
      kernel(output, hidden_states, batch)
  
      assert allclose(output, reference_result)
  
file_path: deepspeed/tests/unit/inference/v2/kernels/ragged_ops/test_logits_gather.py
imports: from typing import List
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.kernels.ragged_ops import RaggedLogitsGather
  from ....v2.inference_test_utils import allclose, get_dtypes
  from .ragged_testing_utils import build_simple_batch
label_desc: Use image node_id 4 for calling a global function with example usage: test_problem_size_permutations(model_dim) without return types