code: def Transformer(
      input_vocab_size,
      output_vocab_size=None,
      d_model=D_MODEL,
      d_ff=D_FF,
      n_encoder_layers=N_LAYERS,
      n_decoder_layers=N_LAYERS,
      n_heads=N_HEADS,
      max_len=MAX_SEQUENCE_LENGTH,
      dropout=DROPOUT_RATE,
      dropout_shared_axes=DROPOUT_SHARED_AXES,
      mode=MODE,
      ff_activation=FF_ACTIVATION_TYPE,
  ):
      """Returns a full Transformer model.
  
      This model is an encoder-decoder that performs tokenized string-to-string
      ("source"-to-"target") transduction:
  
        - inputs (2):
  
            - source: Array representing a batch of text strings via token
              IDs plus padding markers; shape is (batch_size, sequence_length),
              where sequence_length <= ``max_len``. Array elements are integers in
              ``range(input_vocab_size)``, and 0 values mark padding positions.
  
            - target: Array representing a batch of text strings via token
              IDs plus padding markers; shape is (batch_size, sequence_length),
              where sequence_length <= ``max_len``. Array elements are integers in
              ``range(output_vocab_size)``, and 0 values mark padding positions.
  
        - output: 3-D array of raw activations with last/innermost dimension of
          ``output_vocab_size``, suitable for decoding into a batch of token
          strings; shape is (batch_size, sequence_length, ``vocab_size``).
  
      An example use would be to translate (tokenized) sentences from English to
      German.
  
      Args:
        input_vocab_size: Input vocabulary size -- each element of the input tensor
            should be an integer in ``range(vocab_size)``. These integers typically
            represent token IDs from a vocabulary-based tokenizer.
        output_vocab_size: If specified, gives the vocabulary size for the targets;
            if ``None``, then input and target integers (token IDs) are assumed to
            come from the same vocabulary.
        d_model: Last/innermost dimension of activation arrays at most points in
            the model, including the initial embedding output.
        d_ff: Last/innermost dimension of special (typically wider)
            :py:class:`Dense` layer in the feedforward part of each encoder block.
        n_encoder_layers: Number of encoder blocks.
        n_decoder_layers: Number of decoder blocks.
        n_heads: Number of attention heads.
        max_len: Maximum symbol length for positional encoding.
        dropout: Stochastic rate (probability) for dropping an activation value
            when applying dropout within encoder/decoder blocks. The same rate is
            also used for attention dropout in encoder/decoder blocks.
        dropout_shared_axes: Tensor axes on which to share a dropout mask.
            Sharing along batch and sequence axes (``dropout_shared_axes=(0,1)``)
            is a useful way to save memory and apply consistent masks to activation
            vectors at different sequence positions.
        mode: If ``'predict'``, use fast inference. If ``'train'``, each
            encoder/decoder block will include dropout; else, it will pass all
            values through unaltered.
        ff_activation: Type of activation function at the end of each
            encoder/decoder block; must be an activation-type subclass of
            :py:class:`Layer`.
  
      Returns:
        A Transformer model as a layer that maps from a source-target tokenized
        text pair to activations over a vocab set.
      """
      # Avoid 'predict' mode in encoder, since encoder doesn't run stepwise.
      encoder_mode = "eval" if mode == "predict" else mode
  
      # Share embedding weights if no separate output vocab size.
      in_embedder = tl.Embedding(input_vocab_size, d_model)
      if output_vocab_size is None:
          out_embedder = in_embedder
          output_vocab_size = input_vocab_size
      else:
          out_embedder = tl.Embedding(output_vocab_size, d_model)
  
      def _Dropout():
          return tl.Dropout(
              rate=dropout, shared_axes=dropout_shared_axes, mode=mode
          )
  
      def _EncBlock():
          return _EncoderBlock(
              d_model,
              d_ff,
              n_heads,
              dropout,
              dropout_shared_axes,
              mode,
              ff_activation,
          )
  
      def _Encoder():
          encoder = tl.Serial(
              in_embedder,
              _Dropout(),
              tl.PositionalEncoding(max_len=max_len, mode=encoder_mode),
              [_EncBlock() for _ in range(n_encoder_layers)],
              tl.LayerNorm(),
          )
          return tl.Cache(encoder) if mode == "predict" else encoder
  
      def _EncDecBlock():
          return _EncoderDecoderBlock(
              d_model,
              d_ff,
              n_heads,
              dropout,
              dropout_shared_axes,
              mode,
              ff_activation,
          )
  
      # Input to model is encoder-side tokens and decoder-side tokens: tok_d, tok_e
      # Model output is decoder-side vectors and decoder-side tokens: vec_d  tok_d
      return tl.Serial(
          tl.Select(
              [0, 1, 1]
          ),  # Copies decoder tokens for use in loss.
          # Encode.
          tl.Branch([], tl.PaddingMask()),  # tok_e masks tok_d tok_d
          _Encoder(),
          # Decode.
          tl.Select(
              [2, 1, 0]
          ),  # Re-orders inputs: tok_d masks vec_e .....
          tl.ShiftRight(mode=mode),
          out_embedder,
          _Dropout(),
          tl.PositionalEncoding(max_len=max_len, mode=mode),
          tl.Branch(
              [], tl.EncoderDecoderMask()
          ),  # vec_d masks ..... .....
          [_EncDecBlock() for _ in range(n_decoder_layers)],
          tl.LayerNorm(),
          tl.Select([0], n_in=3),  # Drops masks and encoding vectors.
          # Map vectors to match output vocab size.
          tl.Dense(output_vocab_size),
      )
  
file_path: trax/trax/models/transformer.py
imports: from trax import layers
label_desc: Use image node_id 4 for calling a global function with example usage: Transformer(input_vocab_size, output_vocab_size, d_model, d_ff, n_encoder_layers, n_decoder_layers, n_heads, max_len, dropout, dropout_shared_axes, mode, ff_activation) and returns: tl, tl, _EncoderBlock, unknown, _EncoderDecoderBlock