code: def PositionalEncoder(
      mode,
      dropout=None,
      max_len=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 positional encoding layer depending on the arguments.
  
      Args:
        mode: If `'predict'`, use fast inference. If `'train'`, each encoder/decoder
          block will include dropout; else, it will pass all values through
          unaltered.
        dropout: Stochastic rate (probability) for dropping an activation
          value when applying dropout after the embedding block.
        max_len: Maximum symbol length for positional encoding.
        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 layer that will do the positional encoding.
      """
      if not pos_type:
          positional_encoding = tl.PositionalEncoding(
              max_len=max_len,
              dropout=dropout,
              use_bfloat16=use_bfloat16,
              start_from_zero_prob=pos_start_from_zero_prob,
              max_offset_to_add=pos_max_offset_to_add,
              mode=mode,
          )
      elif pos_type == "sin-cos":
          positional_encoding = tl.SinCosPositionalEncoding(mode=mode)
      elif pos_type == "fixed-base":
          positional_encoding = tl.FixedBasePositionalEncoding(
              mode=mode
          )
      elif pos_type == "infinite":
          positional_encoding = tl.InfinitePositionalEncoding(
              affine=False
          )
      elif pos_type == "infinite-affine":
          positional_encoding = tl.InfinitePositionalEncoding()
      elif pos_type == "time-bin":
          positional_encoding = tl.TimeBinPositionalEncoding()
      elif pos_type == "no":
          positional_encoding = (
              tl.Serial()
          )  # no positional encoding at all
      else:  # TODO(lukaszkaiser): name this type and check for the correct name
          assert pos_d_axial_embs is not None
          positional_encoding = tl.AxialPositionalEncoding(
              shape=pos_axial_shape,
              d_embs=pos_d_axial_embs,
              dropout_broadcast_dims=tuple(
                  range(1, len(pos_axial_shape) + 1)
              ),
              dropout=dropout,
              mode=mode,
          )
  
      return positional_encoding
  
file_path: trax/trax/models/research/configurable_transformer.py
imports: from trax import layers
label_desc: Use image node_id 4 for calling a global function with example usage: PositionalEncoder(mode, dropout, max_len, pos_type, pos_axial_shape, pos_d_axial_embs, pos_start_from_zero_prob, pos_max_offset_to_add, use_bfloat16) and returns: positional_encoding