(input_vocab_size,
                 output_vocab_size=None,
                 d_model=512,
                 d_ff=2048,
                 n_encoder_layers=6,
                 n_decoder_layers=6,
                 n_heads=8,
                 dropout=0.1,
                 dropout_shared_axes=None,
                 max_len=2048,
                 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,
                 encoder_attention_type=tl.Attention,
                 n_encoder_attention_layers=1,
                 decoder_attention_type=tl.CausalAttention,
                 n_decoder_attention_layers=2,
                 pos_type=None,
                 pos_axial_shape=None,
                 pos_d_axial_embs=None)
return tl.Serial(
      # Input: encoder_side_tokens, decoder_side_tokens
      # Copy decoder tokens for use in loss.
      tl.Select([0, 0, 1, 1]),          # tok_e tok_e tok_d tok_d

      # Encode.
      tl.Branch([], tl.PaddingMask()),  # tok_e mask_e tok_e tok_d tok_d
      encoder,                          # vec_e mask_e tok_e tok_d tok_d

      # Simple encoder mask, doesn't contain extra dims.
      tl.Select([2, 0, 2], n_in=3),     #  tok_e vec_e tok_e tok_d tok_d
      tl.Fn('EncoderMask',              # mask_e vec_e tok_e tok_d tok_d
            lambda x: x != 0, n_out=1),

      # Decode.
      tl.Select([3, 1, 0, 2]),          #  tok_d vec_e mask_e tok_e tok_d
      tl.ShiftRight(mode=mode),         # stok_d vec_e mask_e tok_e tok_d
      out_encoder,                      # svec_d vec_e mask_e tok_e tok_d

      # Concat encoder and decoder.
      tl.Select([1, 0]),                # vec_e svec_d mask_e tok_e tok_d
      ConcatWithPadding(mode=mode),     # vec_ed tok_e tok_d

      # Decoder blocks with causal attention
      decoder_blocks,                   # vec_ed tok_e tok_d
      tl.LayerNorm(),                   # vec_ed tok_e tok_d

      # Separate out the encoder part from the concatenated vector.
      tl.Select([0, 1, 2, 2]),                     # vec_ed tok_e tok_d tok_d
      StripFromConcatenateWithPadding(mode=mode),  # vec_d tok_d

      # Map to output vocab.
      tl.Dense(output_vocab_size),      # vec_d tok_d
  )
Returns a Transformer model.

  This model expects an input pair: target, source.

  Args:
    input_vocab_size: int: vocab size of the source.
    output_vocab_size: int (optional): vocab size of the target. If None, the
      source and target are assumed to have the same vocab.
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_encoder_layers: int: number of encoder layers
    n_decoder_layers: int: number of decoder layers
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    dropout_shared_axes: axes on which to share dropout mask
    max_len: int: maximum symbol length for positional encoding
    mode: str: 'train' or 'eval'
    ff_activation: the non-linearity in feed-forward 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; if > 0, we use this many SRU layers instead of feed-forward
    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
    encoder_attention_type: The attention layer to use for the encoder part.
    n_encoder_attention_layers: int, within each encoder block, how many
      attention layers to have.
    decoder_attention_type: The attention layer to use for the
      encoder-decoder attention.
    n_decoder_attention_layers: int, within each decoder block, how many
      attention layers to have.
    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 as a layer that maps from a target, source pair to
    activations over a vocab set.
Returns a Transformer model.
Returns a Transformer model .
def Transformer2(input_vocab_size,
                 output_vocab_size=None,
                 d_model=512,
                 d_ff=2048,
                 n_encoder_layers=6,
                 n_decoder_layers=6,
                 n_heads=8,
                 dropout=0.1,
                 dropout_shared_axes=None,
                 max_len=2048,
                 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,
                 encoder_attention_type=tl.Attention,
                 n_encoder_attention_layers=1,
                 decoder_attention_type=tl.CausalAttention,
                 n_decoder_attention_layers=2,
                 pos_type=None,
                 pos_axial_shape=None,
                 pos_d_axial_embs=None):
  """Returns a Transformer model.

  This model expects an input pair: target, source.

  Args:
    input_vocab_size: int: vocab size of the source.
    output_vocab_size: int (optional): vocab size of the target. If None, the
      source and target are assumed to have the same vocab.
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_encoder_layers: int: number of encoder layers
    n_decoder_layers: int: number of decoder layers
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    dropout_shared_axes: axes on which to share dropout mask
    max_len: int: maximum symbol length for positional encoding
    mode: str: 'train' or 'eval'
    ff_activation: the non-linearity in feed-forward 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; if > 0, we use this many SRU layers instead of feed-forward
    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
    encoder_attention_type: The attention layer to use for the encoder part.
    n_encoder_attention_layers: int, within each encoder block, how many
      attention layers to have.
    decoder_attention_type: The attention layer to use for the
      encoder-decoder attention.
    n_decoder_attention_layers: int, within each decoder block, how many
      attention layers to have.
    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 as a layer that maps from a target, source pair to
    activations over a vocab set.
  """
  in_encoder, out_encoder, output_vocab_size = (
      ct.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 = [
      ct.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,
                      n_encoder_attention_layers)
      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
  decoder_blocks = [
      ct.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, decoder_attention_type,
                      n_decoder_attention_layers)
      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, 0, 1, 1]),          # tok_e tok_e tok_d tok_d

      # Encode.
      tl.Branch([], tl.PaddingMask()),  # tok_e mask_e tok_e tok_d tok_d
      encoder,                          # vec_e mask_e tok_e tok_d tok_d

      # Simple encoder mask, doesn't contain extra dims.
      tl.Select([2, 0, 2], n_in=3),     #  tok_e vec_e tok_e tok_d tok_d
      tl.Fn('EncoderMask',              # mask_e vec_e tok_e tok_d tok_d
            lambda x: x != 0, n_out=1),

      # Decode.
      tl.Select([3, 1, 0, 2]),          #  tok_d vec_e mask_e tok_e tok_d
      tl.ShiftRight(mode=mode),         # stok_d vec_e mask_e tok_e tok_d
      out_encoder,                      # svec_d vec_e mask_e tok_e tok_d

      # Concat encoder and decoder.
      tl.Select([1, 0]),                # vec_e svec_d mask_e tok_e tok_d
      ConcatWithPadding(mode=mode),     # vec_ed tok_e tok_d

      # Decoder blocks with causal attention
      decoder_blocks,                   # vec_ed tok_e tok_d
      tl.LayerNorm(),                   # vec_ed tok_e tok_d

      # Separate out the encoder part from the concatenated vector.
      tl.Select([0, 1, 2, 2]),                     # vec_ed tok_e tok_d tok_d
      StripFromConcatenateWithPadding(mode=mode),  # vec_d tok_d

      # Map to output vocab.
      tl.Dense(output_vocab_size),      # vec_d tok_d
  )
def Transformer2 ( input_vocab_size , output_vocab_size = None , d_model = 512 , d_ff = 2048 , n_encoder_layers = 6 , n_decoder_layers = 6 , n_heads = 8 , dropout = 0.1 , dropout_shared_axes = None , max_len = 2048 , 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 , encoder_attention_type = tl . Attention , n_encoder_attention_layers = 1 , decoder_attention_type = tl . CausalAttention , n_decoder_attention_layers = 2 , pos_type = None , pos_axial_shape = None , pos_d_axial_embs = None ) : in_encoder , out_encoder , output_vocab_size = ( ct . 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 = [ ct . 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 , n_encoder_attention_layers ) 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 decoder_blocks = [ ct . 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 , decoder_attention_type , n_decoder_attention_layers ) 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 , 0 , 1 , 1 ] ) , # tok_e tok_e tok_d tok_d # Encode. tl . Branch ( [ ] , tl . PaddingMask ( ) ) , # tok_e mask_e tok_e tok_d tok_d encoder , # vec_e mask_e tok_e tok_d tok_d # Simple encoder mask, doesn't contain extra dims. tl . Select ( [ 2 , 0 , 2 ] , n_in = 3 ) , #  tok_e vec_e tok_e tok_d tok_d tl . Fn ( 'EncoderMask' , # mask_e vec_e tok_e tok_d tok_d lambda x : x != 0 , n_out = 1 ) , # Decode. tl . Select ( [ 3 , 1 , 0 , 2 ] ) , #  tok_d vec_e mask_e tok_e tok_d tl . ShiftRight ( mode = mode ) , # stok_d vec_e mask_e tok_e tok_d out_encoder , # svec_d vec_e mask_e tok_e tok_d # Concat encoder and decoder. tl . Select ( [ 1 , 0 ] ) , # vec_e svec_d mask_e tok_e tok_d ConcatWithPadding ( mode = mode ) , # vec_ed tok_e tok_d # Decoder blocks with causal attention decoder_blocks , # vec_ed tok_e tok_d tl . LayerNorm ( ) , # vec_ed tok_e tok_d # Separate out the encoder part from the concatenated vector. tl . Select ( [ 0 , 1 , 2 , 2 ] ) , # vec_ed tok_e tok_d tok_d StripFromConcatenateWithPadding ( mode = mode ) , # vec_d tok_d # Map to output vocab. tl . Dense ( output_vocab_size ) , # vec_d tok_d )