def EmbeddingAndPositionalEncodings(input_vocab_size,
                                    d_model,
                                    mode,
                                    embedding_dropout,
                                    dropout_shared_axes,
                                    max_len,
                                    output_vocab_size=None,
                                    pos_type=None,
                                    pos_axial_shape=None,
                                    pos_d_axial_embs=None,
                                    pos_start_from_zero_prob=1.0,
                                    pos_max_offset_to_add=0,
                                    use_bfloat16=False):
  """Returns the embedder and positional encoder.

  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.
    d_model: Final dimension of tensors at most points in the model, including
      the initial embedding output.
    mode: If `'predict'`, use fast inference. If `'train'`, each encoder/decoder
      block will include dropout; else, it will pass all values through
      unaltered.
    embedding_dropout: Stochastic rate (probability) for dropping an activation
      value when applying dropout after the embedding 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.
    max_len: Maximum symbol length for positional encoding.
    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.
    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.
    pos_start_from_zero_prob: how often to start from 0 during training,
          (if 1.0, we always start from position 0, if less, we randomize).
    pos_max_offset_to_add: maximum offset to add to positions during training
        when randomizing; this offset plus input length must still be less than
        max_len for all training examples.
    use_bfloat16: If `True`, use bfloat16 weights instead of the default
      float32; this can save memory but may (rarely) lead to numerical issues.

  Returns:
    A tuple of (input encoder, output encoder, output vocab size used).
  """
  # tokens --> vectors
  def Embedder(vocab_size, embedding_mode):
    if vocab_size is not None:
      embedding = tl.Embedding(vocab_size, d_model, use_bfloat16=use_bfloat16)
    else:
      embedding = tl.Dense(d_model, use_bfloat16=use_bfloat16)
    return [
        embedding,
        tl.Dropout(rate=embedding_dropout,
                   shared_axes=dropout_shared_axes,
                   mode=embedding_mode),
    ]

  # NOTE: Positional encodings are not shared between encoder and decoder.

  # Since encoder doesn't run stepwise, we do not use predict mode there.
  encoder_mode = 'eval' if mode == 'predict' else mode
  in_embedder = Embedder(input_vocab_size, encoder_mode)
  in_encoder = in_embedder + [
      PositionalEncoder(encoder_mode,
                        dropout=embedding_dropout,
                        max_len=max_len,
                        pos_type=pos_type,
                        pos_axial_shape=pos_axial_shape,
                        pos_d_axial_embs=pos_d_axial_embs,
                        pos_start_from_zero_prob=pos_start_from_zero_prob,
                        pos_max_offset_to_add=pos_max_offset_to_add,
                        use_bfloat16=use_bfloat16)
  ]

  # If output_vocab_size is None, we reuse the same embedding matrix, otherwise
  # we initialize one.
  assert input_vocab_size or output_vocab_size
  if output_vocab_size is None:
    out_embedder = in_embedder
  else:
    out_embedder = Embedder(output_vocab_size, mode)

  out_encoder = out_embedder + [
      PositionalEncoder(mode,
                        dropout=embedding_dropout,
                        max_len=max_len,
                        pos_type=pos_type,
                        pos_axial_shape=pos_axial_shape,
                        pos_d_axial_embs=pos_d_axial_embs,
                        pos_start_from_zero_prob=pos_start_from_zero_prob,
                        pos_max_offset_to_add=pos_max_offset_to_add,
                        use_bfloat16=use_bfloat16)
  ]

  # Set this to the value actually used.
  if output_vocab_size is None:
    output_vocab_size = input_vocab_size

  if input_vocab_size is None:
    in_encoder = tl.AssertFunction('...a->...b', in_encoder)
  else:
    in_encoder = tl.AssertFunction('...->...d', in_encoder)
  out_encoder = tl.AssertFunction('...->...d', out_encoder)

  return in_encoder, out_encoder, output_vocab_size
Returns the embedder and positional encoder.

  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.
    d_model: Final dimension of tensors at most points in the model, including
      the initial embedding output.
    mode: If `'predict'`, use fast inference. If `'train'`, each encoder/decoder
      block will include dropout; else, it will pass all values through
      unaltered.
    embedding_dropout: Stochastic rate (probability) for dropping an activation
      value when applying dropout after the embedding 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.
    max_len: Maximum symbol length for positional encoding.
    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.
    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.
    pos_start_from_zero_prob: how often to start from 0 during training,
          (if 1.0, we always start from position 0, if less, we randomize).
    pos_max_offset_to_add: maximum offset to add to positions during training
        when randomizing; this offset plus input length must still be less than
        max_len for all training examples.
    use_bfloat16: If `True`, use bfloat16 weights instead of the default
      float32; this can save memory but may (rarely) lead to numerical issues.

  Returns:
    A tuple of (input encoder, output encoder, output vocab size used).
  
trax/models/research/configurable_transformer.py