code: def ConfigurableTransformer(
      input_vocab_size,
      output_vocab_size=None,
      d_model=512,
      d_ff=2048,
      n_encoder_layers=6,
      n_decoder_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,
      encoder_attention_type=tl.Attention,
      encoder_decoder_attention_type=tl.CausalAttention,
      pos_type=None,
      pos_axial_shape=None,
      pos_d_axial_embs=None,
      enc_dec_attention_sparsity=0,
  ):
      """Returns a full Transformer model.
  
      This model is an encoder-decoder that performs tokenized string-to-string
      ("source"-to-"target") transduction:
  
        - inputs (2):
  
            - source: 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(input_vocab_size)`, and `0`
              values mark padding positions.
  
            - target: 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(output_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`).
  
      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: 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
          and decoder 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 an encoder/decoder 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/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 `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: str, 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
        encoder_attention_type: The attention layer to use for the encoder part.
        encoder_decoder_attention_type: The attention layer to use for the
          encoder-decoder attention.
        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.
        enc_dec_attention_sparsity: int, if > 0 use this sparsity in attention.
  
      Returns:
        A Transformer model as a layer that maps from a source-target tokenized
        text pair to activations over a vocab set.
      """
      (
          in_encoder,
          out_encoder,
          output_vocab_size,
      ) = EmbeddingAndPositionalEncodings(
          input_vocab_size,
          d_model,
          mode,
          dropout,
          dropout_shared_axes,
          max_len,
          output_vocab_size=output_vocab_size,
          pos_type=pos_type,
          pos_axial_shape=pos_axial_shape,
          pos_d_axial_embs=pos_d_axial_embs,
      )
  
      # 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,
              encoder_attention_type,
          )
          for i in range(n_encoder_layers)
      ]
      # pylint: enable=g-complex-comprehension
  
      encoder = tl.Serial(in_encoder, encoder_blocks, tl.LayerNorm())
      if mode == "predict":
          encoder = tl.Cache(encoder)
  
      # pylint: disable=g-complex-comprehension
      encoder_decoder_blocks = [
          EncoderDecoderBlock(
              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,
              encoder_decoder_attention_type,
              enc_dec_attention_sparsity,
          )
          for i in range(n_decoder_layers)
      ]
      # pylint: enable=g-complex-comprehension
  
      # 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.SparseDenseWithOptions(  # vec_d .....
              output_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 8 for calling a global function with example usage: ConfigurableTransformer(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, 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, encoder_attention_type, encoder_decoder_attention_type, pos_type, pos_axial_shape, pos_d_axial_embs, enc_dec_attention_sparsity) and returns: tl