code: 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 .....
      )
  
file_path: trax/trax/models/research/rezero.py
imports: from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: ReZeroTransformer(input_vocab_size, output_vocab_size, d_model, d_ff, n_encoder_layers, n_decoder_layers, n_heads, dropout, dropout_shared_axes, max_len, mode, ff_activation) and returns: tl, list