code: def TransformerEncoder(
      vocab_size,
      n_classes=10,
      d_model=D_MODEL,
      d_ff=D_FF,
      n_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 Transformer encoder suitable for N-way classification.
  
      This model maps tokenized text to N-way (``n_classes``) activations:
  
        - input: 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(vocab_size)``, and 0 values mark padding positions.
  
        - output: Array representing a batch of raw (non-normalized) activations
          over ``n_classes`` categories; shape is (batch_size, ``n_classes``).
  
      Args:
        vocab_size: Input vocabulary size -- each element of the input array
            should be an integer in ``range(vocab_size)``. These integers typically
            represent token IDs from a vocabulary-based tokenizer.
        n_classes: Last/innermost dimension of output arrays, suitable for N-way
            classification.
        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_layers: Number of encoder blocks. Each block includes attention, dropout,
            residual, layer-norm, feedforward (:py:class:`Dense`), and activation
            layers.
        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 blocks. The same rate is also
            used for attention dropout in encoder 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 ``'train'``, each encoder block will include dropout; else, it
            will pass all values through unaltered.
        ff_activation: Type of activation function at the end of each encoder
            block; must be an activation-type subclass of :py:class:`Layer`.
  
      Returns:
        A Transformer model that maps strings (conveyed by token IDs) to
        raw (non-normalized) activations over a range of output classes.
      """
  
      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,
          )
  
      return tl.Serial(
          tl.Branch(
              [], tl.PaddingMask()
          ),  # Creates masks from copy of the tokens.
          tl.Embedding(vocab_size, d_model),
          _Dropout(),
          tl.PositionalEncoding(max_len=max_len),
          [_EncBlock() for _ in range(n_layers)],
          tl.Select([0], n_in=2),  # Drops the masks.
          tl.LayerNorm(),
          tl.Mean(axis=1),
          tl.Dense(n_classes),
      )
  
file_path: trax/trax/models/transformer.py
imports: from trax import layers
label_desc: Use image node_id 1 for calling a global function with example usage: TransformerEncoder(vocab_size, n_classes, d_model, d_ff, n_layers, n_heads, max_len, dropout, dropout_shared_axes, mode, ff_activation) and returns: tl, tl, _EncoderBlock