code: def TransformerDecoder(
      vocab_size=None,
      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 decoder.
  
      This model maps sequential inputs to sequential outputs:
  
        - input if ``vocab_size`` is specified: array 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.
  
        - input if ``vocab_size`` is ``None``: 3-D array representing a batch of
          sequences of activation vectors; shape is (batch_size, sequence_length,
          ``d_model``).
  
        - output: 3-D array with shape (batch_size, sequence_length, ``d_model``).
  
      The model uses causal attention and does *not* shift the input to the right.
      Thus, the output for position `t` is based on inputs up to and including
      position `t`.
  
      Args:
        vocab_size: If specified, gives the input vocabulary size -- each element
            of the input tensor should be an integer in ``range(vocab_size)``.
            If ``None``, indicates that the model expects as input sequences of
            floating point vectors, each with ``d_model`` components.
        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 decoder 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 decoder blocks. The same rate is also
            used for attention dropout in 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 ``'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:
        If ``vocab_size`` is defined: a Transformer model that maps strings
        (conveyed by token IDs) to sequences of activation vectors.
  
        If ``vocab_size`` is ``None``: a Transformer model that maps sequences of
        activation vectors to sequences of activation vectors.
      """
  
      def _EmbeddingOrDense():
          return (
              tl.Embedding(vocab_size, d_model)
              if vocab_size is not None
              else tl.Dense(d_model)
          )
  
      def _Dropout():
          return tl.Dropout(
              rate=dropout, shared_axes=dropout_shared_axes, mode=mode
          )
  
      def _DecBlock():
          return _DecoderBlock(
              d_model,
              d_ff,
              n_heads,
              dropout,
              dropout_shared_axes,
              mode,
              ff_activation,
          )
  
      return tl.Serial(
          _EmbeddingOrDense(),
          _Dropout(),
          tl.PositionalEncoding(max_len=max_len),
          [_DecBlock() for _ in range(n_layers)],
          tl.LayerNorm(),
      )
  
file_path: trax/trax/models/transformer.py
imports: from trax import layers
label_desc: Use image node_id 2 for calling a global function with example usage: TransformerDecoder(vocab_size, d_model, d_ff, n_layers, n_heads, max_len, dropout, dropout_shared_axes, mode, ff_activation) and returns: tl, unknown, tl, _DecoderBlock