code: def _mask_probs(probs, eos_token, finished):
      """Masks log probabilities.
  
      The result is that finished beams allocate all probability mass to eos and
      unfinished beams remain unchanged.
  
      Args:
        probs: Log probabilities of shape `[batch_size, beam_width, vocab_size]`
        eos_token: An int32 id corresponding to the EOS token to allocate
          probability to.
        finished: A boolean tensor of shape `[batch_size, beam_width]` that
          specifies which elements in the beam are finished already.
  
      Returns:
        A tensor of shape `[batch_size, beam_width, vocab_size]`, where
        unfinished beams stay unchanged and finished beams are replaced with a
        tensor with all probability on the EOS token.
      """
      vocab_size = tf.shape(probs)[2]
      # All finished examples are replaced with a vector that has all
      # probability on EOS
      finished_row = tf.one_hot(
          eos_token,
          vocab_size,
          dtype=probs.dtype,
          on_value=tf.convert_to_tensor(0.0, dtype=probs.dtype),
          off_value=probs.dtype.min,
      )
      finished_probs = tf.tile(
          tf.reshape(finished_row, [1, 1, -1]),
          tf.concat([tf.shape(finished), [1]], 0),
      )
      finished_mask = tf.tile(
          tf.expand_dims(finished, 2), [1, 1, vocab_size]
      )
  
      return tf.where(finished_mask, finished_probs, probs)
  
imports: import collections
  import numpy
  import tensorflow
  from tensorflow_addons import options
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import decoder
  from tensorflow_addons.utils import keras_utils
  from tensorflow_addons.utils.resource_loader import LazySO
  from tensorflow_addons.utils.types import FloatTensorLike, TensorLike, Number
  from typeguard import typechecked
  from typing import Callable, Optional
label_desc: Use image node_id 15 for calling a global function with example usage: _mask_probs(probs, eos_token, finished) and returns: tf