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(),
  )
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.
  
trax/models/transformer.py