(d_model, d_ff, n_heads, dropout, ff_activation,
                        ff_dropout, mode, ff_use_sru=0, ff_chunk_size=0,
                        ff_sparsity=0)
return [                             # vec_d1 vec_d2 vec_e masks
      causal_attention_half_residual,
      tl.ReversibleSwap(),
      enc_dec_attention_half_residual,
      tl.ReversibleSwap(),
      tl.ReversibleHalfResidual(feed_forward),
      tl.ReversibleSwap(),
  ]
Reversible transformer decoder layer.

  Args:
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    ff_activation: the non-linearity in feed-forward layer
    ff_dropout: float: (optional) separate dropout rate for feed-forward layer
    mode: str: 'train' or 'eval'
    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

  Returns:
    the layer.
Reversible transformer decoder layer.
Reversible transformer decoder layer .
def EncoderDecoderBlock(d_model, d_ff, n_heads, dropout, ff_activation,
                        ff_dropout, mode, ff_use_sru=0, ff_chunk_size=0,
                        ff_sparsity=0):
  """Reversible transformer decoder layer.

  Args:
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    ff_activation: the non-linearity in feed-forward layer
    ff_dropout: float: (optional) separate dropout rate for feed-forward layer
    mode: str: 'train' or 'eval'
    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

  Returns:
    the layer.
  """
  enc_dec_attention = tl.EncDecAttention(
      n_heads=n_heads, d_qk=d_model//n_heads, d_v=d_model//n_heads,
      attention_dropout=dropout, output_dropout=dropout,
      mode=mode)
  enc_dec_attention_half_residual = tl.ReversibleHalfResidual(
      tl.LayerNorm(),
      attention_layer=enc_dec_attention,
  )

  causal_attention = tl.SelfAttention(
      n_heads=n_heads, d_qk=d_model//n_heads, d_v=d_model//n_heads,
      causal=True,
      attention_dropout=dropout, output_dropout=dropout,
      mode=mode)
  causal_attention_half_residual = tl.ReversibleHalfResidual(
      tl.LayerNorm(),
      attention_layer=causal_attention,
  )

  feed_forward = ct.FeedForwardWithOptions(
      d_model, d_ff, dropout, [-2], ff_activation, ff_dropout,
      ff_chunk_size, ff_use_sru, ff_sparsity, True, mode)

  return [                             # vec_d1 vec_d2 vec_e masks
      causal_attention_half_residual,
      tl.ReversibleSwap(),
      enc_dec_attention_half_residual,
      tl.ReversibleSwap(),
      tl.ReversibleHalfResidual(feed_forward),
      tl.ReversibleSwap(),
  ]
def EncoderDecoderBlock ( d_model , d_ff , n_heads , dropout , ff_activation , ff_dropout , mode , ff_use_sru = 0 , ff_chunk_size = 0 , ff_sparsity = 0 ) : enc_dec_attention = tl . EncDecAttention ( n_heads = n_heads , d_qk = d_model // n_heads , d_v = d_model // n_heads , attention_dropout = dropout , output_dropout = dropout , mode = mode ) enc_dec_attention_half_residual = tl . ReversibleHalfResidual ( tl . LayerNorm ( ) , attention_layer = enc_dec_attention , ) causal_attention = tl . SelfAttention ( n_heads = n_heads , d_qk = d_model // n_heads , d_v = d_model // n_heads , causal = True , attention_dropout = dropout , output_dropout = dropout , mode = mode ) causal_attention_half_residual = tl . ReversibleHalfResidual ( tl . LayerNorm ( ) , attention_layer = causal_attention , ) feed_forward = ct . FeedForwardWithOptions ( d_model , d_ff , dropout , [ - 2 ] , ff_activation , ff_dropout , ff_chunk_size , ff_use_sru , ff_sparsity , True , mode ) return [ # vec_d1 vec_d2 vec_e masks causal_attention_half_residual , tl . ReversibleSwap ( ) , enc_dec_attention_half_residual , tl . ReversibleSwap ( ) , tl . ReversibleHalfResidual ( feed_forward ) , tl . ReversibleSwap ( ) , ]