code: def test_update_scores(use_unk_dist):
      vocab_size = 10
      batch_beam_size = 3
      us = sockeye.beam_search.UpdateScores(prevent_unk=use_unk_dist)
      pad_dist = onp.full(
          (1, vocab_size), fill_value=onp.inf, dtype="float32"
      )
      pad_dist[0, 0] = 0
      eos_dist = onp.full(
          (batch_beam_size, vocab_size),
          fill_value=onp.inf,
          dtype="float32",
      )
      eos_dist[:, C.EOS_ID] = 0
  
      lengths = onp.array([0, 1, 0], dtype="int32")
      max_lengths = onp.array(
          [1, 2, 3], dtype="int32"
      )  # first on reaches max length
      scores_accumulated = onp.ones((3, 1), dtype="float32")
      finished = pt.tensor(
          [
              False,  # not finished
              True,  # finished
              False,
          ],  # not finished
          dtype=pt.bool,
      )
      target_dists = onp.random.uniform(0, 1, (3, vocab_size)).astype(
          "float32"
      )
  
      scores, lengths = us(
          pt.tensor(target_dists),
          finished,
          pt.tensor(scores_accumulated),
          pt.tensor(lengths),
          pt.tensor(max_lengths),
          pt.tensor(pad_dist),
          pt.tensor(eos_dist),
      )
      scores = scores.detach().numpy()
      lengths = lengths
      pt.testing.assert_close(
          lengths, pt.tensor([1, 1, 1], dtype=pt.int32)
      )  # all lengths but finished updated + 1
      assert (
          scores[0] == (1.0 + target_dists[0] + eos_dist)
      ).all()  # 1 reached max length, force eos
      assert (
          scores[1] == (1.0 + pad_dist[0]).tolist()
      ).all()  # 2 finished, force pad, keep score
      if use_unk_dist:
          assert (
              scores[2, C.UNK_ID] == onp.inf
          )  # 3 scores of <unk> should be np.inf
          target_dists[2, C.UNK_ID] = onp.inf
          assert (
              scores[2] == (1.0 + target_dists[2])
          ).all()  # 3 scores + previous scores
      else:
          assert (scores[2] == (1.0 + target_dists[2])).all()
  
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 12 for calling a global function with example usage: test_update_scores(use_unk_dist) without return types