(input_vocab_size,
                      output_vocab_size=None,
                      d_model=512,
                      d_ff=2048,
                      n_encoder_layers=6,
                      n_decoder_layers=6,
                      n_heads=8,
                      dropout=0.1,
                      dropout_shared_axes=None,
                      max_len=2048,
                      mode='train',
                      ff_activation=tl.Relu)
return tl.Serial(
      # Input: encoder_side_tokens, decoder_side_tokens
      # Copy decoder tokens for use in loss.
      tl.Select([0, 1, 1]),               # tok_e tok_d tok_d

      # Encode.
      tl.Branch([], tl.PaddingMask()),    # tok_e masks ..... .....
      encoder,                            # vec_e ..... ..... .....

      # Decode.
      tl.Select([2, 1, 0]),               # tok_d masks vec_e .....
      tl.ShiftRight(mode=mode),           # tok_d ..... ..... .....
      out_encoder,                        # vec_d ..... ..... .....
      tl.Branch(
          [], tl.EncoderDecoderMask()),   # vec_d masks ..... .....
      encoder_decoder_blocks,             # vec_d masks ..... .....
      tl.LayerNorm(),                     # vec_d ..... ..... .....

      # Map to output vocab.
      tl.Select([0], n_in=3),             # vec_d tok_d
      tl.Dense(output_vocab_size),        # vec_d .....
  )
Returns a ReZero transformer model.

  This model expects an input pair: source, target.

  Args:
    input_vocab_size: int: vocab size of the source.
    output_vocab_size: int (optional): vocab size of the target. If None, the
      source and target are assumed to have the same vocab.
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_encoder_layers: int: number of encoder layers
    n_decoder_layers: int: number of decoder layers
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    dropout_shared_axes: axes on which to share dropout mask
    max_len: int: maximum symbol length for positional encoding
    mode: str: 'train' or 'eval'
    ff_activation: the non-linearity in feed-forward layer

  Returns:
    A ReZero transformer model as a layer that maps from a source, target pair
    to activations over a vocab set.
Returns a ReZero transformer model.
Returns a ReZero transformer model .
def ReZeroTransformer(input_vocab_size,
                      output_vocab_size=None,
                      d_model=512,
                      d_ff=2048,
                      n_encoder_layers=6,
                      n_decoder_layers=6,
                      n_heads=8,
                      dropout=0.1,
                      dropout_shared_axes=None,
                      max_len=2048,
                      mode='train',
                      ff_activation=tl.Relu):
  """Returns a ReZero transformer model.

  This model expects an input pair: source, target.

  Args:
    input_vocab_size: int: vocab size of the source.
    output_vocab_size: int (optional): vocab size of the target. If None, the
      source and target are assumed to have the same vocab.
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_encoder_layers: int: number of encoder layers
    n_decoder_layers: int: number of decoder layers
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    dropout_shared_axes: axes on which to share dropout mask
    max_len: int: maximum symbol length for positional encoding
    mode: str: 'train' or 'eval'
    ff_activation: the non-linearity in feed-forward layer

  Returns:
    A ReZero transformer model as a layer that maps from a source, target pair
    to activations over a vocab set.
  """
  def Embedder(vocab_size):  # tokens --> vectors
    return [
        tl.Embedding(vocab_size, d_model),
        tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
    ]

  in_embedder = Embedder(input_vocab_size)
  out_embedder = (in_embedder if output_vocab_size is None
                  else Embedder(output_vocab_size))

  # Positional encoding are not shared between encoder and decoder.
  # Since encoder doesn't run stepwise, we do not use predict mode there.
  encoder_mode = 'eval' if mode == 'predict' else mode
  in_encoder = in_embedder + [
      tl.PositionalEncoding(max_len=max_len, mode=encoder_mode)
  ]
  out_encoder = out_embedder + [
      tl.PositionalEncoding(max_len=max_len, mode=mode)
  ]

  if output_vocab_size is None:
    output_vocab_size = input_vocab_size

  encoder_blocks = [
      _EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
                    mode, ff_activation)
      for i in range(n_encoder_layers)]

  encoder = tl.Serial(
      in_encoder,
      encoder_blocks,
      tl.LayerNorm()
  )
  if mode == 'predict':
    encoder = tl.Cache(encoder)

  encoder_decoder_blocks = [
      _EncoderDecoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
                           mode, ff_activation)
      for i in range(n_decoder_layers)]

  # Assemble and return the model.
  return tl.Serial(
      # Input: encoder_side_tokens, decoder_side_tokens
      # Copy decoder tokens for use in loss.
      tl.Select([0, 1, 1]),               # tok_e tok_d tok_d

      # Encode.
      tl.Branch([], tl.PaddingMask()),    # tok_e masks ..... .....
      encoder,                            # vec_e ..... ..... .....

      # Decode.
      tl.Select([2, 1, 0]),               # tok_d masks vec_e .....
      tl.ShiftRight(mode=mode),           # tok_d ..... ..... .....
      out_encoder,                        # vec_d ..... ..... .....
      tl.Branch(
          [], tl.EncoderDecoderMask()),   # vec_d masks ..... .....
      encoder_decoder_blocks,             # vec_d masks ..... .....
      tl.LayerNorm(),                     # vec_d ..... ..... .....

      # Map to output vocab.
      tl.Select([0], n_in=3),             # vec_d tok_d
      tl.Dense(output_vocab_size),        # vec_d .....
  )
def ReZeroTransformer ( input_vocab_size , output_vocab_size = None , d_model = 512 , d_ff = 2048 , n_encoder_layers = 6 , n_decoder_layers = 6 , n_heads = 8 , dropout = 0.1 , dropout_shared_axes = None , max_len = 2048 , mode = 'train' , ff_activation = tl . Relu ) : def Embedder ( vocab_size ) : # tokens --> vectors return [ tl . Embedding ( vocab_size , d_model ) , tl . Dropout ( rate = dropout , shared_axes = dropout_shared_axes , mode = mode ) , ] in_embedder = Embedder ( input_vocab_size ) out_embedder = ( in_embedder if output_vocab_size is None else Embedder ( output_vocab_size ) ) # Positional encoding are not shared between encoder and decoder. # Since encoder doesn't run stepwise, we do not use predict mode there. encoder_mode = 'eval' if mode == 'predict' else mode in_encoder = in_embedder + [ tl . PositionalEncoding ( max_len = max_len , mode = encoder_mode ) ] out_encoder = out_embedder + [ tl . PositionalEncoding ( max_len = max_len , mode = mode ) ] if output_vocab_size is None : output_vocab_size = input_vocab_size encoder_blocks = [ _EncoderBlock ( d_model , d_ff , n_heads , dropout , dropout_shared_axes , mode , ff_activation ) for i in range ( n_encoder_layers ) ] encoder = tl . Serial ( in_encoder , encoder_blocks , tl . LayerNorm ( ) ) if mode == 'predict' : encoder = tl . Cache ( encoder ) encoder_decoder_blocks = [ _EncoderDecoderBlock ( d_model , d_ff , n_heads , dropout , dropout_shared_axes , mode , ff_activation ) for i in range ( n_decoder_layers ) ] # Assemble and return the model. return tl . Serial ( # Input: encoder_side_tokens, decoder_side_tokens # Copy decoder tokens for use in loss. tl . Select ( [ 0 , 1 , 1 ] ) , # tok_e tok_d tok_d # Encode. tl . Branch ( [ ] , tl . PaddingMask ( ) ) , # tok_e masks ..... ..... encoder , # vec_e ..... ..... ..... # Decode. tl . Select ( [ 2 , 1 , 0 ] ) , # tok_d masks vec_e ..... tl . ShiftRight ( mode = mode ) , # tok_d ..... ..... ..... out_encoder , # vec_d ..... ..... ..... tl . Branch ( [ ] , tl . EncoderDecoderMask ( ) ) , # vec_d masks ..... ..... encoder_decoder_blocks , # vec_d masks ..... ..... tl . LayerNorm ( ) , # vec_d ..... ..... ..... # Map to output vocab. tl . Select ( [ 0 ] , n_in = 3 ) , # vec_d tok_d tl . Dense ( output_vocab_size ) , # vec_d ..... )