code: 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
  
file_path: trax/trax/models/research/configurable_transformer.py
imports: from trax import layers
label_desc: Use image node_id 5 for calling a global function with example usage: EmbeddingAndPositionalEncodings(input_vocab_size, d_model, mode, embedding_dropout, dropout_shared_axes, max_len, output_vocab_size, pos_type, pos_axial_shape, pos_d_axial_embs, pos_start_from_zero_prob, pos_max_offset_to_add, use_bfloat16) and returns: in_encoder,  out_encoder,  output_vocab_size, list