(attention_type, d_model, d_ff, n_heads, dropout,
                          dropout_shared_axes, mode, ff_activation,
                          context_bias_layer, location_bias_layer,
                          total_pooling)
return [
      tl.Residual(  # vecs
          tl.LayerNorm(),
          attention,
          _Dropout(),
      ),  # vecs
      tl.Residual(
          tl.LayerNorm(),
          feed_forward,
          _Dropout(),
      ),  # vecs
  ]
Returns a list of layers.

    The layers implement a Transformer decoder block with relative attention
  parametrization.

  The input to the block is a pair, (activations, mask), where the mask was
  created from the original source tokens to prevent attending to the padding
  part of the input.

  Args:
    attention_type: attention type.
    d_model: Final dimension of tensors at most points in the model, including
      the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each block.
    n_heads: Number of attention heads.
    dropout: Stochastic rate (probability) for dropping an activation value when
      applying dropout within a block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
      along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
      way to save memory and apply consistent masks to activation vectors at
      different sequence positions.
    mode: If `'train'`, each block will include dropout; else, it will pass all
      values through unaltered.
    ff_activation: Type of activation function at the end of each block; must be
      an activation-type subclass of `Layer`.
    context_bias_layer: context bias layer.
    location_bias_layer: location bias layer.
    total_pooling: The combined pool size of previously used funnel blocks.

  Returns:
    A list of layers that maps (activations, att_vecs, mask) to
                               (activations, att_vecs, mask).
Returns a list of layers.
Returns a list of layers .
def _RelativeDecoderBlock(attention_type, d_model, d_ff, n_heads, dropout,
                          dropout_shared_axes, mode, ff_activation,
                          context_bias_layer, location_bias_layer,
                          total_pooling):
  """Returns a list of layers.

    The layers implement a Transformer decoder block with relative attention
  parametrization.

  The input to the block is a pair, (activations, mask), where the mask was
  created from the original source tokens to prevent attending to the padding
  part of the input.

  Args:
    attention_type: attention type.
    d_model: Final dimension of tensors at most points in the model, including
      the initial embedding output.
    d_ff: Size of special dense layer in the feed-forward part of each block.
    n_heads: Number of attention heads.
    dropout: Stochastic rate (probability) for dropping an activation value when
      applying dropout within a block.
    dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
      along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
      way to save memory and apply consistent masks to activation vectors at
      different sequence positions.
    mode: If `'train'`, each block will include dropout; else, it will pass all
      values through unaltered.
    ff_activation: Type of activation function at the end of each block; must be
      an activation-type subclass of `Layer`.
    context_bias_layer: context bias layer.
    location_bias_layer: location bias layer.
    total_pooling: The combined pool size of previously used funnel blocks.

  Returns:
    A list of layers that maps (activations, att_vecs, mask) to
                               (activations, att_vecs, mask).
  """
  if attention_type == RelativeAttentionWrapper:
    attention = RelativeAttentionWrapper(
        d_model,
        n_heads,
        dropout,
        mode=mode,
        context_bias_layer=context_bias_layer,
        location_bias_layer=location_bias_layer,
        total_pooling=total_pooling)
  else:
    attention = ApplyAttentionLayer(
        attention_type,
        d_model,
        n_heads,
        d_model // n_heads,
        d_model // n_heads,
        causal=True,
        masked=False,
        attention_dropout=dropout,
        output_dropout=dropout,
        attention_chunk_size=0,  # Disables tl.Chunk in ApplyAttentionLayer.
        mode=mode,
    )

  feed_forward = FeedForwardBlock(d_model, d_ff, dropout, dropout_shared_axes,
                                  mode, ff_activation)

  def _Dropout():
    return tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)

  return [
      tl.Residual(  # vecs
          tl.LayerNorm(),
          attention,
          _Dropout(),
      ),  # vecs
      tl.Residual(
          tl.LayerNorm(),
          feed_forward,
          _Dropout(),
      ),  # vecs
  ]
def _RelativeDecoderBlock ( attention_type , d_model , d_ff , n_heads , dropout , dropout_shared_axes , mode , ff_activation , context_bias_layer , location_bias_layer , total_pooling ) : if attention_type == RelativeAttentionWrapper : attention = RelativeAttentionWrapper ( d_model , n_heads , dropout , mode = mode , context_bias_layer = context_bias_layer , location_bias_layer = location_bias_layer , total_pooling = total_pooling ) else : attention = ApplyAttentionLayer ( attention_type , d_model , n_heads , d_model // n_heads , d_model // n_heads , causal = True , masked = False , attention_dropout = dropout , output_dropout = dropout , attention_chunk_size = 0 , # Disables tl.Chunk in ApplyAttentionLayer. mode = mode , ) feed_forward = FeedForwardBlock ( d_model , d_ff , dropout , dropout_shared_axes , mode , ff_activation ) def _Dropout ( ) : return tl . Dropout ( rate = dropout , shared_axes = dropout_shared_axes , mode = mode ) return [ tl . Residual ( # vecs tl . LayerNorm ( ) , attention , _Dropout ( ) , ) , # vecs tl . Residual ( tl . LayerNorm ( ) , feed_forward , _Dropout ( ) , ) , # vecs ]