code: def TransformerLM(vocab_size,
                    feature_depth=512,
                    feedforward_depth=2048,
                    num_layers=6,
                    num_heads=8,
                    dropout=0.1,
                    max_len=2048,
                    mode='train'):
    
    return layers.Serial(
        layers.ShiftRight(),
        layers.Embedding(feature_depth, vocab_size),
        layers.Dropout(rate=dropout, mode=mode),
        layers.PositionalEncoding(max_len=max_len),
        layers.Serial(*[DecoderLayer(feature_depth, feedforward_depth, num_heads,
                                     dropout, mode)
                        for _ in range(num_layers)]),
        layers.LayerNorm(),
        layers.Dense(vocab_size),
        layers.LogSoftmax()
    )
docs: Transformer language model (only uses the decoder part of Transformer).
  
    Args:
      vocab_size: int: vocab size
      feature_depth: int:  depth of embedding
      feedforward_depth: int: depth of feed-forward layer
      num_layers: int: number of encoder/decoder layers
      num_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'
  
    Returns:
      the layer.