code: def ReZeroTransformerLM(
      vocab_size,
      d_model=512,
      d_ff=2048,
      n_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 language model.
  
      The input to the model is a tensor of tokens. (This model uses only the
      decoder part of the overall ReZero 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)
        dropout_shared_axes: axes on which to share dropout mask
        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 ReZero transformer language model as a layer that maps from a tensor of
        tokens to activations over a vocab set.
      """
      positional_encoder = [
          tl.Embedding(vocab_size, d_model),
          tl.Dropout(
              rate=dropout, shared_axes=dropout_shared_axes, 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,
              dropout_shared_axes,
              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
      )
  
file_path: trax/trax/models/research/rezero.py
imports: from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: ReZeroTransformerLM(vocab_size, d_model, d_ff, n_layers, n_heads, dropout, dropout_shared_axes, max_len, mode, ff_activation) and returns: tl