original_docstring: Returns a Transformer language model.
  
    The input to the model is a tensor of tokens. (This model uses only the
    decoder part of the overall Transformer.)
  
    Args:
      vocab_size: int: vocab size
      d_model: int:  depth of embedding
      d_ff: int: depth of feed-forward layer
      n_layers: int: number of encoder/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', 'eval' or 'predict', predict mode is for fast inference
      ff_activation: the non-linearity in feed-forward layer
  
    Returns:
      A Transformer language model as a layer that maps from a tensor of tokens
      to activations over a vocab set.
    
docstring: Returns a Transformer language model.
  The input to the model is a tensor of tokens. (This model uses only the
  decoder part of the overall Transformer.)
docstring_tokens: Returns a Transformer language model . The input to the model is a tensor of tokens . ( This model uses only the decoder part of the overall Transformer . )
code: def TransformerLM(vocab_size,
                    d_model=512,
                    d_ff=2048,
                    n_layers=6,
                    n_heads=8,
                    dropout=0.1,
                    max_len=2048,
                    mode='train',
                    ff_activation=tl.Relu):
    """Returns a Transformer language model.
  
    The input to the model is a tensor of tokens. (This model uses only the
    decoder part of the overall Transformer.)
  
    Args:
      vocab_size: int: vocab size
      d_model: int:  depth of embedding
      d_ff: int: depth of feed-forward layer
      n_layers: int: number of encoder/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', 'eval' or 'predict', predict mode is for fast inference
      ff_activation: the non-linearity in feed-forward layer
  
    Returns:
      A Transformer language model as a layer that maps from a tensor of tokens
      to activations over a vocab set.
    """
    positional_encoder = [
        tl.Embedding(d_model, vocab_size),
        tl.Dropout(rate=dropout, name='embedding', mode=mode),
        tl.PositionalEncoding(max_len=max_len, mode=mode)]
  
    decoder_blocks = [
        # pylint: disable=g-complex-comprehension
        _DecoderBlock(d_model, d_ff, n_heads,
                      dropout, i, mode, ff_activation)
        for i in range(n_layers)]
  
    # Assemble and return the model.
    return tl.Serial(              # tokens (or chunked tuple of tokens)
        tl.ShiftRight(mode=mode),  # toks
        positional_encoder,        # vecs
        decoder_blocks,            # vecs
        tl.LayerNorm(),            # vecs
        tl.Dense(vocab_size),      # vecs
        tl.LogSoftmax(),           # vecs
    )
code_tokens: def TransformerLM ( vocab_size , d_model = 512 , d_ff = 2048 , n_layers = 6 , n_heads = 8 , dropout = 0.1 , max_len = 2048 , mode = 'train' , ff_activation = tl . Relu ) : positional_encoder = [ tl . Embedding ( d_model , vocab_size ) , tl . Dropout ( rate = dropout , name = 'embedding' , mode = mode ) , tl . PositionalEncoding ( max_len = max_len , mode = mode ) ] decoder_blocks = [ _DecoderBlock ( d_model , d_ff , n_heads , dropout , i , mode , ff_activation ) for i in range ( n_layers ) ] return tl . Serial ( tl . ShiftRight ( mode = mode ) , positional_encoder , decoder_blocks , tl . LayerNorm ( ) , tl . Dense ( vocab_size ) , tl . LogSoftmax ( ) , )
short_docstring: Returns a Transformer language model.
short_docstring_tokens: Returns a Transformer language model .
comment: """Returns a Transformer language model.
  
    The input to the model is a tensor of tokens. (This model uses only the
    decoder part of the overall Transformer.)
  
    Args:
      vocab_size: int: vocab size
      d_model: int:  depth of embedding
      d_ff: int: depth of feed-forward layer
      n_layers: int: number of encoder/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', 'eval' or 'predict', predict mode is for fast inference
      ff_activation: the non-linearity in feed-forward layer
  
    Returns:
      A Transformer language model as a layer that maps from a tensor of tokens
      to activations over a vocab set.
    """$SEP$# pylint: disable=g-complex-comprehension$SEP$# Assemble and return the model.$SEP$# tokens (or chunked tuple of tokens)$SEP$# toks$SEP$# vecs$SEP$# vecs$SEP$# vecs$SEP$# vecs$SEP$# vecs