code: def ConfigurableTransformerLM(
      vocab_size,
      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",
      loss_sparsity_type="mult",
      loss_sparsity=0,
      loss_d_lowrank=0,
      loss_sparsity_prob=None,
      attention_chunk_size=0,
      attention_type=tl.CausalAttention,
      pos_type=None,
      pos_axial_shape=None,
      pos_d_axial_embs=None,
      pos_start_from_zero_prob=1.0,
      pos_max_offset_to_add=0,
  ):
      """Returns a Transformer language model.
  
      This model performs autoregressive language modeling:
  
        - 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 3 tensor representing a batch of log-probability
          distributions for each sequence position over possible token IDs;
          shape is (batch_size, sequence_length, `vocab_size`).
  
      This model uses only the decoder part of the overall Transformer.
  
      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.
        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 `'predict'`, use fast inference. 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'`
        loss_sparsity_type: string, type of sparsity to used in loss layer. See
          SparseDenseWithOptions for options. None if no sparsity should be used.
        loss_sparsity: int, the sparsity for loss layer (if used)
        loss_d_lowrank: int, the dimensions for intermediate layer (if used)
        loss_sparsity_prob: float, the probability for sparse version of loss to be
          used. If None, only sparse version is used.
        attention_chunk_size: int, if > 0 run attention chunked at this size
        attention_type: The attention layer to use for the decoder 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.
        pos_start_from_zero_prob: how often to start from 0 during training,
          (if 1.0, we always start from position 0, if less, we randomize).
        pos_max_offset_to_add: maximum offset to add to positions during training
          when randomizing; this offset plus input length must still be less than
          max_len for all training examples.
  
      Returns:
        A 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
          ),
          PositionalEncoder(
              mode,
              dropout,
              max_len,
              pos_type,
              pos_axial_shape,
              pos_d_axial_embs,
              pos_start_from_zero_prob,
              pos_max_offset_to_add,
          ),
      ]
  
      # pylint: disable=g-complex-comprehension
      decoder_blocks = [
          DecoderBlock(
              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(  # tokens (or chunked tuple of tokens)
          tl.ShiftRight(mode=mode),  # toks
          positional_encoder,  # vecs
          decoder_blocks,  # vecs
          tl.LayerNorm(),  # vecs
          tl.SparseDenseWithOptions(  # vecs
              vocab_size,
              d_input=d_model,
              sparsity_type=loss_sparsity_type,
              sparsity=loss_sparsity,
              d_lowrank=loss_d_lowrank,
              prob_sparse=loss_sparsity_prob,
              mode=mode,
          ),
      )
  
file_path: trax/trax/models/research/configurable_transformer.py
imports: from trax import layers
label_desc: Use image node_id 7 for calling a global function with example usage: ConfigurableTransformerLM(vocab_size, 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, loss_sparsity_type, loss_sparsity, loss_d_lowrank, loss_sparsity_prob, attention_chunk_size, attention_type, pos_type, pos_axial_shape, pos_d_axial_embs, pos_start_from_zero_prob, pos_max_offset_to_add) and returns: tl