code: def _ragged_embed_test_helper(
      sequence_config: List[Tuple[int, int]],
      embed_dtype: torch.dtype,
      token_dtype: torch.dtype,
      embed_dim: int,
      vocab_size: int,
      do_padding: bool = False,
      pos_embed_size: int = -1,
      pos_embed_offset: int = 0,
  ) -> None:
      """
      Helper for embedding test to limit the number of tests to run.
  
      Params:
          embed_dim (int): Model dimension
          vocab_size (int): Leading dimension on embedding weight
          pos_embed_size (int): Size of positional embedding. If negative, no positional embedding
              is used.
          pos_embed_offset (int): Offset for positional embedding. Effectively, the raw offsets
              of a token into a sequence are offset by this amount into the embedding matrix. (
              i.e. the shape of the positional embeddings is (pos_embed_size + pos_embed_offset
              embed_dim)
      """
      device = get_accelerator().current_device()
  
      # Heads/Block size are irrelevant here but need something.
      (
          batch,
          _,
          _,
      ) = build_batch_and_manager(
          sequence_config,
          64,
          16,
          64,
          vocab_range=vocab_size,
          padding=do_padding,
      )
  
      embedding_table = torch.randn(
          (vocab_size, embed_dim), dtype=embed_dtype, device=device
      )
  
      if pos_embed_size > 0:
          pos_embedding_table = torch.randn(
              (pos_embed_size + pos_embed_offset, embed_dim),
              dtype=embed_dtype,
              device=device,
          )
          positional_ids = (
              torch.cat(
                  [
                      torch.arange(
                          start_idx,
                          start_idx + seq_len,
                          dtype=token_dtype,
                          device=device,
                      )
                      for seq_len, start_idx in sequence_config
                  ]
              )
              + pos_embed_offset
          )
      else:
          pos_embedding_table = None
          positional_ids = None
  
      baseline_output = baseline_implementation(
          batch.input_ids().to(token_dtype),
          embedding_table,
          batch.current_tokens,
          pos_embedding_table,
          positional_ids,
      )
  
      kernel = RaggedEmbeddingKernel(
          embed_dtype, token_dtype, embed_dim
      )
      output = torch.empty_like(baseline_output)
  
      kernel(
          output,
          batch,
          embedding_table,
          position_embed_weight=pos_embedding_table,
          position_embed_offset=pos_embed_offset,
      )
  
      if do_padding:
          assert output.shape[0] != batch.current_tokens
  
      assert allclose(
          output[: batch.current_tokens],
          baseline_output[: batch.current_tokens],
      )
  
file_path: DeepSpeed/tests/unit/inference/v2/kernels/ragged_ops/test_ragged_embed.py
imports: from typing import List, Optional, Tuple
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.kernels.ragged_ops import RaggedEmbeddingKernel
  from ....v2.inference_test_utils import allclose, get_dtypes
  from .ragged_testing_utils import build_batch_and_manager
label_desc: Use image node_id 2 for calling a global function with example usage: _ragged_embed_test_helper(sequence_config, embed_dtype, token_dtype, embed_dim, vocab_size, do_padding, pos_embed_size, pos_embed_offset) without return types