code: def test_beam_search():
      device = pt.device("cpu")
      dtype = pt.float32
      num_source_factors = 1
      num_target_factors = 1
      vocab_size = len(C.VOCAB_SYMBOLS) + 1  # 1 actual word: word id 4
      beam_size = 1
      bos_id = 2
      eos_id = 3
  
      inference = _TestInference(output_vocab_size=vocab_size)
      bs = sockeye.beam_search.BeamSearch(
          beam_size=beam_size,
          dtype=dtype,
          bos_id=bos_id,
          eos_id=eos_id,
          device=device,
          output_vocab_size=vocab_size,
          scorer=sockeye.beam_search.CandidateScorer(),
          num_source_factors=num_source_factors,
          num_target_factors=num_target_factors,
          inference=inference,
          beam_search_stop=C.BEAM_SEARCH_STOP_ALL,
          sample=None,
      )
  
      # inputs
      batch_size = 1
      max_length = 3
      source = pt.tensor(
          [[C.BOS_ID, 4, C.EOS_ID, C.PAD_ID, C.PAD_ID]], dtype=dtype
      ).reshape(1, -1, 1)
      source_length = (
          (source != C.PAD_ID).sum(1).reshape(-1)
      )  # (batch_size,)
  
      restrict_lexicon = None
      max_output_lengths = pt.tensor([max_length], dtype=pt.int)
  
      bs_out = bs(
          source, source_length, restrict_lexicon, max_output_lengths
      )
      r = bs_out
  
      print("beam search lengths", r.lengths)
      print("internal lengths", inference.states[0])
      pt.testing.assert_close(r.lengths, inference.states[0].squeeze(1))
      assert inference.states[1] == max_length
  
imports: from typing import List, Optional
  from typing import Tuple
  import numpy
  import pytest
  import torch
  import numpy
  import sockeye.beam_search
  import sockeye.constants
  import sockeye.lexicon
  import sockeye.utils
label_desc: Use image node_id 13 for calling a global function with example usage: test_beam_search() without return types