(d_model, d_ff, d_attention_key, d_attention_value,
                 n_heads, attention_type, dropout, ff_activation,
                 ff_dropout, ff_use_sru, ff_chunk_size, ff_sparsity,
                 attention_chunk_size, n_attention_layers=1,
                 n_feedforward_layers=1, center_layernorm=True,
                 use_bfloat16=False, mode='train')
return ([_attention_half_residual() for _ in range(n_attention_layers)]
          + [_feed_forward() for _ in range(n_feedforward_layers)])
Reversible transformer decoder layer.

  Args:
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    d_attention_key: int: depth of key vector for each attention head
    d_attention_value: int: depth of value vector for each attention head
    n_heads: int: number of attention heads
    attention_type: subclass of tl.BaseCausalAttention: attention class to use
    dropout: float: dropout rate (how much to drop out)
    ff_activation: the non-linearity in feed-forward layer
    ff_dropout: the dropout rate in feed-forward layer
    ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
    ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
    ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
    attention_chunk_size: int, if > 0 run attention chunked at this size
    n_attention_layers: how many residual causal attention layers should we
      have before the feed-forward block (default: 1, the standard block)
    n_feedforward_layers: how many FFNN layers should we have (default 1).
    center_layernorm: whether to use centering in LayerNorm (default) or if
      to skip it, which is known as RMS normalization.
    use_bfloat16: whether to use bfloat16 for weights (default: False).
    mode: str: 'train' or 'eval'


  Returns:
    the layer.
Reversible transformer decoder layer.
Reversible transformer decoder layer .
def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value,
                 n_heads, attention_type, dropout, ff_activation,
                 ff_dropout, ff_use_sru, ff_chunk_size, ff_sparsity,
                 attention_chunk_size, n_attention_layers=1,
                 n_feedforward_layers=1, center_layernorm=True,
                 use_bfloat16=False, mode='train'):
  """Reversible transformer decoder layer.

  Args:
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    d_attention_key: int: depth of key vector for each attention head
    d_attention_value: int: depth of value vector for each attention head
    n_heads: int: number of attention heads
    attention_type: subclass of tl.BaseCausalAttention: attention class to use
    dropout: float: dropout rate (how much to drop out)
    ff_activation: the non-linearity in feed-forward layer
    ff_dropout: the dropout rate in feed-forward layer
    ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
    ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
    ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
    attention_chunk_size: int, if > 0 run attention chunked at this size
    n_attention_layers: how many residual causal attention layers should we
      have before the feed-forward block (default: 1, the standard block)
    n_feedforward_layers: how many FFNN layers should we have (default 1).
    center_layernorm: whether to use centering in LayerNorm (default) or if
      to skip it, which is known as RMS normalization.
    use_bfloat16: whether to use bfloat16 for weights (default: False).
    mode: str: 'train' or 'eval'


  Returns:
    the layer.
  """
  # pylint: disable=g-complex-comprehension
  def _Attn():
    return ct.ApplyAttentionLayer(
        attention_type, d_model, n_heads, d_attention_key,
        d_attention_value, True, False, dropout, dropout,
        attention_chunk_size, mode)

  def _FF():
    return ct.FeedForwardWithOptions(
        d_model, d_ff, dropout, [-2], ff_activation, ff_dropout,
        ff_chunk_size, ff_use_sru, ff_sparsity, center_layernorm,
        mode, use_bfloat16)

  def _attention_half_residual():
    return [
        tl.ReversibleHalfResidual(tl.LayerNorm(center=center_layernorm),
                                  attention_layer=_Attn(),
                                  name='ReversibleHalfResidualDecoderAttn'),
        tl.ReversibleSwap()
    ]

  def _feed_forward():
    return [
        tl.ReversibleHalfResidual(_FF(),
                                  name='ReversibleHalfResidualDecoderFF'),
        tl.ReversibleSwap()
    ]

  return ([_attention_half_residual() for _ in range(n_attention_layers)]
          + [_feed_forward() for _ in range(n_feedforward_layers)])
def DecoderBlock ( d_model , d_ff , d_attention_key , d_attention_value , n_heads , attention_type , dropout , ff_activation , ff_dropout , ff_use_sru , ff_chunk_size , ff_sparsity , attention_chunk_size , n_attention_layers = 1 , n_feedforward_layers = 1 , center_layernorm = True , use_bfloat16 = False , mode = 'train' ) : # pylint: disable=g-complex-comprehension def _Attn ( ) : return ct . ApplyAttentionLayer ( attention_type , d_model , n_heads , d_attention_key , d_attention_value , True , False , dropout , dropout , attention_chunk_size , mode ) def _FF ( ) : return ct . FeedForwardWithOptions ( d_model , d_ff , dropout , [ - 2 ] , ff_activation , ff_dropout , ff_chunk_size , ff_use_sru , ff_sparsity , center_layernorm , mode , use_bfloat16 ) def _attention_half_residual ( ) : return [ tl . ReversibleHalfResidual ( tl . LayerNorm ( center = center_layernorm ) , attention_layer = _Attn ( ) , name = 'ReversibleHalfResidualDecoderAttn' ) , tl . ReversibleSwap ( ) ] def _feed_forward ( ) : return [ tl . ReversibleHalfResidual ( _FF ( ) , name = 'ReversibleHalfResidualDecoderFF' ) , tl . ReversibleSwap ( ) ] return ( [ _attention_half_residual ( ) for _ in range ( n_attention_layers ) ] + [ _feed_forward ( ) for _ in range ( n_feedforward_layers ) ] )