original_docstring: Returns a Transformer model.
  
    This model expects an input pair: target, source.
  
    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)
      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 Transformer model as a layer that maps from a target, source pair to
      activations over a vocab set.
    
docstring: Returns a Transformer model.
  This model expects an input pair: target, source.
docstring_tokens: Returns a Transformer model . This model expects an input pair : target source .
code: def Transformer(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,
                  max_len=2048,
                  mode='train',
                  ff_activation=tl.Relu):
    """Returns a Transformer model.
  
    This model expects an input pair: target, source.
  
    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)
      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 Transformer model as a layer that maps from a target, source pair to
      activations over a vocab set.
    """
    def PositionalEncoder(vocab_size):  # tokens --> vectors
      return [
          tl.Embedding(d_model, vocab_size),
          tl.Dropout(rate=dropout, mode=mode),
          tl.PositionalEncoding(max_len=max_len),
      ]
  
    in_encoder = PositionalEncoder(input_vocab_size)
    out_encoder = (in_encoder if output_vocab_size is None
                   else PositionalEncoder(output_vocab_size))
    if output_vocab_size is None:
      output_vocab_size = input_vocab_size
  
    encoder_blocks = [
        _EncoderBlock(
            d_model, d_ff, n_heads, dropout, i, 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, i, 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(),                    # 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 .....
        tl.LogSoftmax(),                    # vec_d .....
    )
code_tokens: def Transformer ( 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 , max_len = 2048 , mode = 'train' , ff_activation = tl . Relu ) : def PositionalEncoder ( vocab_size ) : return [ tl . Embedding ( d_model , vocab_size ) , tl . Dropout ( rate = dropout , mode = mode ) , tl . PositionalEncoding ( max_len = max_len ) , ] in_encoder = PositionalEncoder ( input_vocab_size ) out_encoder = ( in_encoder if output_vocab_size is None else PositionalEncoder ( output_vocab_size ) ) if output_vocab_size is None : output_vocab_size = input_vocab_size encoder_blocks = [ _EncoderBlock ( d_model , d_ff , n_heads , dropout , i , 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 , i , mode , ff_activation ) for i in range ( n_decoder_layers ) ] return tl . Serial ( tl . Select ( [ 0 , 1 , 1 ] ) , tl . Branch ( [ ] , tl . PaddingMask ( ) ) , encoder , tl . Select ( [ 2 , 1 , 0 ] ) , tl . ShiftRight ( ) , out_encoder , tl . Branch ( [ ] , tl . EncoderDecoderMask ( ) ) , encoder_decoder_blocks , tl . LayerNorm ( ) , tl . Select ( [ 0 ] , n_in = 3 ) , tl . Dense ( output_vocab_size ) , tl . LogSoftmax ( ) , )
short_docstring: Returns a Transformer model.
short_docstring_tokens: Returns a Transformer model .
comment: """Returns a Transformer model.
  
    This model expects an input pair: target, source.
  
    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)
      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 Transformer model as a layer that maps from a target, source pair to
      activations over a vocab set.
    """$SEP$# tokens --> vectors$SEP$# Assemble and return the model.$SEP$# Input: encoder_side_tokens, decoder_side_tokens$SEP$# Copy decoder tokens for use in loss.$SEP$# tok_e tok_d tok_d$SEP$# Encode.$SEP$# tok_e masks ..... .....$SEP$# vec_e ..... ..... .....$SEP$# Decode.$SEP$# tok_d masks vec_e .....$SEP$# tok_d ..... ..... .....$SEP$# vec_d ..... ..... .....$SEP$# vec_d masks ..... .....$SEP$# vec_d masks ..... .....$SEP$# vec_d ..... ..... .....$SEP$# Map to output vocab.$SEP$# vec_d tok_d$SEP$# vec_d .....$SEP$# vec_d .....