original_docstring: Returns a Transformer encoder model.
  
    The input to the model is a tensor of tokens.
  
    Args:
      vocab_size: int: vocab size
      n_classes: how many classes on output
      d_model: int:  depth of embedding
      d_ff: int: depth of feed-forward layer
      n_layers: int: number of encoder/decoder layers
      n_heads: int: number of attention heads
      dropout: float: dropout rate (how much to drop out)
      max_len: int: maximum symbol length for positional encoding
      mode: str: 'train' or 'eval'
      ff_activation: the non-linearity in feed-forward layer
  
    Returns:
      A Transformer model as a layer that maps from a tensor of tokens to
      activations over a set of output classes.
    
docstring: Returns a Transformer encoder model.
  The input to the model is a tensor of tokens.
docstring_tokens: Returns a Transformer encoder model . The input to the model is a tensor of tokens .
code: def TransformerEncoder(vocab_size,
                         n_classes=10,
                         d_model=512,
                         d_ff=2048,
                         n_layers=6,
                         n_heads=8,
                         dropout=0.1,
                         max_len=2048,
                         mode='train',
                         ff_activation=tl.Relu):
    """Returns a Transformer encoder model.
  
    The input to the model is a tensor of tokens.
  
    Args:
      vocab_size: int: vocab size
      n_classes: how many classes on output
      d_model: int:  depth of embedding
      d_ff: int: depth of feed-forward layer
      n_layers: int: number of encoder/decoder layers
      n_heads: int: number of attention heads
      dropout: float: dropout rate (how much to drop out)
      max_len: int: maximum symbol length for positional encoding
      mode: str: 'train' or 'eval'
      ff_activation: the non-linearity in feed-forward layer
  
    Returns:
      A Transformer model as a layer that maps from a tensor of tokens to
      activations over a set of output classes.
    """
    positional_encoder = [
        tl.Embedding(d_model, vocab_size),
        tl.Dropout(rate=dropout, name='emb_dropout', mode=mode),
        tl.PositionalEncoding(max_len=max_len)]
  
    encoder_blocks = [
        _EncoderBlock(d_model, d_ff, n_heads, dropout, i, mode, ff_activation)
        for i in range(n_layers)]
  
    # Assemble and return the model.
    return tl.Serial(                               # toks
        # Encode.
        tl.Branch(
            positional_encoder, tl.PaddingMask()),  # vecs masks
        encoder_blocks,                             # vecs masks
        tl.Select([0], n_in=2),                     # vecs
        tl.LayerNorm(),                             # vecs
  
        # Map to output categories.
        tl.Mean(axis=1),                            # vecs
        tl.Dense(n_classes),                        # vecs
        tl.LogSoftmax(),                            # vecs
    )
code_tokens: def TransformerEncoder ( vocab_size , n_classes = 10 , d_model = 512 , d_ff = 2048 , n_layers = 6 , n_heads = 8 , dropout = 0.1 , max_len = 2048 , mode = 'train' , ff_activation = tl . Relu ) : positional_encoder = [ tl . Embedding ( d_model , vocab_size ) , tl . Dropout ( rate = dropout , name = 'emb_dropout' , mode = mode ) , tl . PositionalEncoding ( max_len = max_len ) ] encoder_blocks = [ _EncoderBlock ( d_model , d_ff , n_heads , dropout , i , mode , ff_activation ) for i in range ( n_layers ) ] return tl . Serial ( tl . Branch ( positional_encoder , tl . PaddingMask ( ) ) , encoder_blocks , tl . Select ( [ 0 ] , n_in = 2 ) , tl . LayerNorm ( ) , tl . Mean ( axis = 1 ) , tl . Dense ( n_classes ) , tl . LogSoftmax ( ) , )
short_docstring: Returns a Transformer encoder model.
short_docstring_tokens: Returns a Transformer encoder model .
comment: """Returns a Transformer encoder model.
  
    The input to the model is a tensor of tokens.
  
    Args:
      vocab_size: int: vocab size
      n_classes: how many classes on output
      d_model: int:  depth of embedding
      d_ff: int: depth of feed-forward layer
      n_layers: int: number of encoder/decoder layers
      n_heads: int: number of attention heads
      dropout: float: dropout rate (how much to drop out)
      max_len: int: maximum symbol length for positional encoding
      mode: str: 'train' or 'eval'
      ff_activation: the non-linearity in feed-forward layer
  
    Returns:
      A Transformer model as a layer that maps from a tensor of tokens to
      activations over a set of output classes.
    """$SEP$# Assemble and return the model.$SEP$# toks$SEP$# Encode.$SEP$# vecs masks$SEP$# vecs masks$SEP$# vecs$SEP$# vecs$SEP$# Map to output categories.$SEP$# vecs$SEP$# vecs$SEP$# vecs