code: def test_get_best_word_indices_for_kth_hypotheses():
      # data
      all_hyp_indices = np.array(
          [
              [0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 2, 0, 0, 4, 3],
              [0, 2, 2, 0, 1, 0, 0, 2, 1, 1, 3, 1, 1, 0, 1, 4, 0, 4],
              [0, 1, 0, 1, 2, 1, 4, 3, 2, 3, 0, 4, 3, 1, 2, 1, 1, 0],
              [0, 1, 0, 0, 3, 2, 2, 1, 3, 4, 4, 2, 2, 3, 3, 2, 2, 1],
              [0, 2, 4, 1, 4, 2, 3, 4, 4, 2, 0, 3, 4, 4, 4, 3, 3, 2],
          ],
          dtype="int32",
      )
      ks = [
          np.array([0]),
          np.array([1]),
          np.array([2]),
          np.array([3]),
          np.array([4]),
      ]
      expected_indices = [
          np.array(
              [[2, 1, 0, 0, 0, 0, 1, 3, 3, 2, 0, 0, 0, 1, 1, 2, 3]],
              dtype="int32",
          ),
          np.array(
              [[1, 2, 1, 2, 2, 3, 4, 4, 4, 3, 1, 1, 1, 2, 2, 3, 4]],
              dtype="int32",
          ),
          np.array(
              [[2, 1, 0, 0, 0, 1, 0, 0, 0, 0, 4, 2, 3, 3, 3, 4, 0]],
              dtype="int32",
          ),
          np.array(
              [[2, 1, 0, 0, 0, 1, 0, 0, 0, 0, 2, 3, 2, 0, 0, 0, 1]],
              dtype="int32",
          ),
          np.array(
              [[2, 1, 0, 1, 1, 2, 3, 2, 2, 4, 3, 4, 4, 4, 4, 1, 2]],
              dtype="int32",
          ),
      ]
  
      # extract individually
      for k, expected_result in zip(ks, expected_indices):
          result = sockeye.inference.Translator._get_best_word_indices_for_kth_hypotheses(
              k, all_hyp_indices
          )
          assert result.shape == expected_result.shape
          assert (result == expected_result).all()
  
      # extract all at once
      ks = np.concatenate(ks, axis=0)
      expected_indices = np.concatenate(expected_indices, axis=0)
      result = sockeye.inference.Translator._get_best_word_indices_for_kth_hypotheses(
          ks, all_hyp_indices
      )
      assert result.shape == expected_indices.shape
      assert (result == expected_indices).all()
  
imports: import itertools
  import json
  from math import ceil
  from unittest.mock import patch, Mock
  import numpy
  import pytest
  import torch
  import sockeye.beam_search
  import sockeye.constants
  import sockeye.data_io
  import sockeye.inference
  import sockeye.lexicon
  import sockeye.model
  import sockeye.utils
label_desc: Use image node_id 15 for calling a global function with example usage: test_get_best_word_indices_for_kth_hypotheses() without return types