class_docstr: Beam search decoder.
  
  **NOTE** If you are using the `BeamSearchDecoder` with a cell wrapped in
  `tfa.seq2seq.AttentionWrapper`, then you must ensure that:
  
  - The encoder output has been tiled to `beam_width` via
    `tfa.seq2seq.tile_batch` (NOT `tf.tile`).
  - The `batch_size` argument passed to the `get_initial_state` method of
    this wrapper is equal to `true_batch_size * beam_width`.
  - The initial state created with `get_initial_state` above contains a
    `cell_state` value containing properly tiled final state from the
    encoder.
  
  An example:
  
  ```
  tiled_encoder_outputs = tfa.seq2seq.tile_batch(
      encoder_outputs, multiplier=beam_width)
  tiled_encoder_final_state = tfa.seq2seq.tile_batch(
      encoder_final_state, multiplier=beam_width)
  tiled_sequence_length = tfa.seq2seq.tile_batch(
      sequence_length, multiplier=beam_width)
  attention_mechanism = MyFavoriteAttentionMechanism(
      num_units=attention_depth,
      memory=tiled_inputs,
      memory_sequence_length=tiled_sequence_length)
  attention_cell = AttentionWrapper(cell, attention_mechanism, ...)
  decoder_initial_state = attention_cell.get_initial_state(
      batch_size=true_batch_size * beam_width, dtype=dtype)
  decoder_initial_state = decoder_initial_state.clone(
      cell_state=tiled_encoder_final_state)
  ```
  
  Meanwhile, with `tfa.seq2seq.AttentionWrapper`, coverage penalty is suggested to use
  when computing scores (https://arxiv.org/pdf/1609.08144.pdf). It encourages
  the decoding to cover all inputs.
docstr: Initialize the decoder.
  
  Args:
    embedding: A `Tensor` (or `Variable`) to pass as the `params` argument
      for `tf.nn.embedding_lookup`. This overrides `embedding_fn` set in
      the constructor.
    start_tokens: Start the decoding from these tokens.
      A `int32` `Tensor` of shape `[batch_size]`.
    end_token: The token that marks the end of decoding.
      A `int32` scalar `Tensor`.
    initial_state: The initial cell state as a (possibly nested) structure
      of `Tensor` and `TensorArray`.
  
  Returns:
    `(finished, start_inputs, initial_state)`.
  
  Raises:
    ValueError: If `embedding` is `None` and `embedding_fn` was not set
      in the constructor.
    ValueError: If `start_tokens` is not a vector or `end_token` is not a
      scalar.
docstr_tok: ["Initialize","the","decoder",".","Args",":","embedding",":","A","`","Tensor","`","(","or","`","Variable","`",")","to","pass","as","the","`","params","`","argument","for","`","tf.nn.embedding_lookup","`",".","This","overrides","`","embedding_fn","`","set","in","the","constructor",".","start_tokens",":","Start","the","decoding","from","these","tokens",".","A","`","int32","`","`","Tensor","`","of","shape","`","[","batch_size","]","`",".","end_token",":","The","token","that","marks","the","end","of","decoding",".","A","`","int32","`","scalar","`","Tensor","`",".","initial_state",":","The","initial","cell","state","as","a","(","possibly","nested",")","structure","of","`","Tensor","`","and","`","TensorArray","`",".","Returns",":","`","(","finished",",","start_inputs",",","initial_state",")","`",".","Raises",":","ValueError",":","If","`","embedding","`","is","`","None","`","and","`","embedding_fn","`","was","not","set","in","the","constructor",".","ValueError",":","If","`","start_tokens","`","is","not","a","vector","or","`","end_token","`","is","not","a","scalar","."]
code: def initialize(
      self, embedding, start_tokens, end_token, initial_state
  ):
      """Initialize the decoder.
  
      Args:
        embedding: A `Tensor` (or `Variable`) to pass as the `params` argument
          for `tf.nn.embedding_lookup`. This overrides `embedding_fn` set in
          the constructor.
        start_tokens: Start the decoding from these tokens.
          A `int32` `Tensor` of shape `[batch_size]`.
        end_token: The token that marks the end of decoding.
          A `int32` scalar `Tensor`.
        initial_state: The initial cell state as a (possibly nested) structure
          of `Tensor` and `TensorArray`.
  
      Returns:
        `(finished, start_inputs, initial_state)`.
  
      Raises:
        ValueError: If `embedding` is `None` and `embedding_fn` was not set
          in the constructor.
        ValueError: If `start_tokens` is not a vector or `end_token` is not a
          scalar.
      """
      if embedding is not None:
          self._embedding_fn = lambda ids: tf.nn.embedding_lookup(
              embedding, ids
          )
      elif self._embedding_fn is None:
          raise ValueError(
              "You should either pass an embedding variable when calling the "
              "BeamSearchDecoder or set embedding_fn in the constructor."
          )
  
      self._start_tokens = tf.convert_to_tensor(
          start_tokens, dtype=tf.int32, name="start_tokens"
      )
      if self._start_tokens.shape.ndims != 1:
          raise ValueError("start_tokens must be a vector")
      self._end_token = tf.convert_to_tensor(
          end_token, dtype=tf.int32, name="end_token"
      )
      if self._end_token.shape.ndims != 0:
          raise ValueError("end_token must be a scalar")
  
      self._batch_size = tf.size(start_tokens)
      self._initial_cell_state = tf.nest.map_structure(
          self._maybe_split_batch_beams,
          initial_state,
          self._cell.state_size,
      )
      self._start_tokens = tf.tile(
          tf.expand_dims(self._start_tokens, 1), [1, self._beam_width]
      )
      self._start_inputs = self._embedding_fn(self._start_tokens)
  
      self._finished = tf.one_hot(
          tf.zeros([self._batch_size], dtype=tf.int32),
          depth=self._beam_width,
          on_value=False,
          off_value=True,
          dtype=tf.bool,
      )
  
      finished, start_inputs = self._finished, self._start_inputs
  
      dtype = tf.nest.flatten(self._initial_cell_state)[0].dtype
      log_probs = tf.one_hot(  # shape(batch_sz, beam_sz)
          tf.zeros([self._batch_size], dtype=tf.int32),
          depth=self._beam_width,
          on_value=tf.convert_to_tensor(0.0, dtype=dtype),
          off_value=tf.convert_to_tensor(-np.Inf, dtype=dtype),
          dtype=dtype,
      )
      init_attention_probs = get_attention_probs(
          self._initial_cell_state, self._coverage_penalty_weight
      )
      if init_attention_probs is None:
          init_attention_probs = ()
  
      initial_state = BeamSearchDecoderState(
          cell_state=self._initial_cell_state,
          log_probs=log_probs,
          finished=finished,
          lengths=tf.zeros(
              [self._batch_size, self._beam_width], dtype=tf.int64
          ),
          accumulated_attention_probs=init_attention_probs,
      )
  
      return (finished, start_inputs, initial_state)
  
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 2 for calling the BeamSearchDecoder obj's underlying member method code with example usage: obj.initialize(embedding, start_tokens, end_token, initial_state) and returns: finished,  start_inputs,  initial_state