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 BeamSearchDecoder.
  
  Args:
    cell: A layer that implements the `tf.keras.layers.AbstractRNNCell`
      interface.
    beam_width:  Python integer, the number of beams.
    embedding_fn: A callable that takes a `int32` `Tensor` of token IDs
      and returns embedding tensors. If set, the `embedding` argument in
      the decoder call should be set to `None`.
    output_layer: (Optional) An instance of `tf.keras.layers.Layer`,
      i.e., `tf.keras.layers.Dense`.  Optional layer to apply to the RNN
      output prior to storing the result or sampling.
    length_penalty_weight: Float weight to penalize length. Disabled with
      0.0.
    coverage_penalty_weight: Float weight to penalize the coverage of
      source sentence. Disabled with 0.0.
    reorder_tensor_arrays: If `True`, `TensorArray`s' elements within the
      cell state will be reordered according to the beam search path. If
      the `TensorArray` can be reordered, the stacked form will be
      returned. Otherwise, the `TensorArray` will be returned as is. Set
      this flag to `False` if the cell state contains `TensorArray`s that
      are not amenable to reordering.
    **kwargs: Dict, other keyword arguments for initialization.
docstr_tok: ["Initialize","the","BeamSearchDecoder",".","Args",":","cell",":","A","layer","that","implements","the","`","tf.keras.layers.AbstractRNNCell","`","interface",".","beam_width",":","Python","integer",",","the","number","of","beams",".","embedding_fn",":","A","callable","that","takes","a","`","int32","`","`","Tensor","`","of","token","IDs","and","returns","embedding","tensors",".","If","set",",","the","`","embedding","`","argument","in","the","decoder","call","should","be","set","to","`","None","`",".","output_layer",":","(","Optional",")","An","instance","of","`","tf.keras.layers.Layer","`",",","i.e.",",","`","tf.keras.layers.Dense","`",".","Optional","layer","to","apply","to","the","RNN","output","prior","to","storing","the","result","or","sampling",".","length_penalty_weight",":","Float","weight","to","penalize","length",".","Disabled","with","0.0.","coverage_penalty_weight",":","Float","weight","to","penalize","the","coverage","of","source","sentence",".","Disabled","with","0.0.","reorder_tensor_arrays",":","If","`","True","`",",","`","TensorArray","`","s","'","elements","within","the","cell","state","will","be","reordered","according","to","the","beam","search","path",".","If","the","`","TensorArray","`","can","be","reordered",",","the","stacked","form","will","be","returned",".","Otherwise",",","the","`","TensorArray","`","will","be","returned","as","is",".","Set","this","flag","to","`","False","`","if","the","cell","state","contains","`","TensorArray","`","s","that","are","not","amenable","to","reordering",".","*","*","kwargs",":","Dict",",","other","keyword","arguments","for","initialization","."]
code: def __init__(
      self,
      cell: tf.keras.layers.Layer,
      beam_width: int,
      embedding_fn: Optional[Callable] = None,
      output_layer: Optional[tf.keras.layers.Layer] = None,
      length_penalty_weight: FloatTensorLike = 0.0,
      coverage_penalty_weight: FloatTensorLike = 0.0,
      reorder_tensor_arrays: bool = True,
      **kwargs,
  ):
      """Initialize the BeamSearchDecoder.
  
      Args:
        cell: A layer that implements the `tf.keras.layers.AbstractRNNCell`
          interface.
        beam_width:  Python integer, the number of beams.
        embedding_fn: A callable that takes a `int32` `Tensor` of token IDs
          and returns embedding tensors. If set, the `embedding` argument in
          the decoder call should be set to `None`.
        output_layer: (Optional) An instance of `tf.keras.layers.Layer`,
          i.e., `tf.keras.layers.Dense`.  Optional layer to apply to the RNN
          output prior to storing the result or sampling.
        length_penalty_weight: Float weight to penalize length. Disabled with
          0.0.
        coverage_penalty_weight: Float weight to penalize the coverage of
          source sentence. Disabled with 0.0.
        reorder_tensor_arrays: If `True`, `TensorArray`s' elements within the
          cell state will be reordered according to the beam search path. If
          the `TensorArray` can be reordered, the stacked form will be
          returned. Otherwise, the `TensorArray` will be returned as is. Set
          this flag to `False` if the cell state contains `TensorArray`s that
          are not amenable to reordering.
        **kwargs: Dict, other keyword arguments for initialization.
      """
      super().__init__(
          cell,
          beam_width,
          output_layer=output_layer,
          length_penalty_weight=length_penalty_weight,
          coverage_penalty_weight=coverage_penalty_weight,
          reorder_tensor_arrays=reorder_tensor_arrays,
          **kwargs,
      )
  
      self._embedding_fn = embedding_fn
  
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 1 to create a new BeamSearchDecoder object from inherited base classes: BeamSearchDecoderMixin, decoder with example: obj = BeamSearchDecoder(cell, beam_width, embedding_fn, output_layer, length_penalty_weight, coverage_penalty_weight, reorder_tensor_arrays)