code: def baseline_implementation(
      hidden_states: torch.Tensor, seq_lens: List[int]
  ) -> torch.Tensor:
      output = torch.empty(
          (len(seq_lens), hidden_states.shape[1]),
          dtype=hidden_states.dtype,
          device=hidden_states.device,
      )
  
      offset = 0
      for i, seq_len in enumerate(seq_lens):
          output[i] = hidden_states[offset + seq_len - 1]
          offset += seq_len
  
      return output
  
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 1 for calling a global function with example usage: baseline_implementation(hidden_states, seq_lens) and returns: output