code: def ConfigurableTransformerEncoder(
      vocab_size,
      n_classes=10,
      d_model=512,
      d_ff=2048,
      n_layers=6,
      n_heads=8,
      max_len=2048,
      dropout=0.1,
      dropout_shared_axes=None,
      mode="train",
      ff_activation=tl.Relu,
      ff_dropout=0.1,
      ff_chunk_size=0,
      ff_use_sru=0,
      ff_sparsity=0,
      ff_sparsity_type="1inN",
      attention_chunk_size=0,
      attention_type=tl.Attention,
      pos_type=None,
      pos_axial_shape=None,
      pos_d_axial_embs=None,
  ):
      """Returns a Transformer encoder merged with an N-way categorization head.
  
      This model performs text categorization:
  
        - input: rank 2 tensor representing a batch of text strings via token IDs
          plus padding markers; shape is (batch_size, sequence_length). The tensor
          elements are integers in `range(vocab_size)`, and `0` values mark padding
          positions.
  
        - output: rank 2 tensor representing a batch of log-probability
          distributions over N categories; shape is (batch_size, `n_classes`).
  
      Args:
        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.
        n_classes: Final dimension of the output tensors, representing N-way
          classification.
        d_model: Final dimension of tensors at most points in the model, including
          the initial embedding output.
        d_ff: Size of special dense layer in the feed-forward part of each encoder
          block.
        n_layers: Number of encoder blocks. Each block includes attention, dropout,
          residual, feed-forward (`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 an encoder block.
        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 `Layer`.
        ff_dropout: Stochastic rate (probability) for dropping an activation value
          when applying dropout after the FF dense layer.
        ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
        ff_use_sru: int or pair of ints; if > 0, we use this many SRU layers
          in addition to the feed-forward block (second int specifies sru size)
        ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
        ff_sparsity_type: string, if ff_sparsity >0,
          use SparseFF if ff_sparsity_type=`'1inN'` and
          use BlockSparseFF if ff_sparsity_type=`'Block'`
        attention_chunk_size: int, if > 0 run attention chunked at this size
        attention_type: The attention layer to use for the encoder part.
        pos_type: string, the type of positional embeddings to use.
        pos_axial_shape: tuple of ints: input shape to use for the axial position
          encoding. If unset, axial position encoding is disabled.
        pos_d_axial_embs: tuple of ints: depth of position embedding for each axis.
          Tuple length must match pos_axial_shape, and values must sum to d_model.
  
      Returns:
        A Transformer model that maps strings (conveyed via token IDs) to
        probability-like activations over a range of output classes.
      """
      positional_encoder = [
          tl.Embedding(vocab_size, d_model),
          tl.Dropout(
              rate=dropout, shared_axes=dropout_shared_axes, mode=mode
          ),
          PositionalEncoder(
              mode,
              dropout,
              max_len,
              pos_type,
              pos_axial_shape,
              pos_d_axial_embs,
          ),
      ]
  
      positional_encoder = tl.AssertFunction(
          "...->...d", positional_encoder
      )
  
      # pylint: disable=g-complex-comprehension
      encoder_blocks = [
          EncoderBlock(
              d_model,
              d_ff,
              n_heads,
              dropout,
              dropout_shared_axes,
              mode,
              ff_activation,
              ff_dropout,
              ff_chunk_size,
              ff_use_sru,
              ff_sparsity,
              ff_sparsity_type,
              attention_chunk_size,
              attention_type,
          )
          for i in range(n_layers)
      ]
      # pylint: enable=g-complex-comprehension
  
      # Assemble and return the model.
      return tl.Serial(  # toks
          # Encode.
          tl.Branch(positional_encoder, tl.PaddingMask()),  # vecs masks
          encoder_blocks,  # vecs masks
          tl.Select([0], n_in=2),  # vecs
          tl.LayerNorm(),  # vecs
          # Map to output categories.
          tl.Mean(axis=1),  # vecs
          tl.Dense(n_classes),  # vecs
      )
  
file_path: trax/trax/models/research/configurable_transformer.py
imports: from trax import layers
label_desc: Use image node_id 6 for calling a global function with example usage: ConfigurableTransformerEncoder(vocab_size, n_classes, d_model, d_ff, n_layers, n_heads, max_len, dropout, dropout_shared_axes, mode, ff_activation, ff_dropout, ff_chunk_size, ff_use_sru, ff_sparsity, ff_sparsity_type, attention_chunk_size, attention_type, pos_type, pos_axial_shape, pos_d_axial_embs) and returns: tl