def __init__(self, max_len=2048, dropout=0.0, dropout_broadcast_dims=(-2,),
               use_bfloat16=False, start_from_zero_prob=1.0,
               max_offset_to_add=0, d_feature=None, mode='train'):
    """Creates a :py:class:`PositionalEncoding` instance in a given mode.

    Args:
      max_len: Maximum input sequence length.
      dropout: Probability of *not* adding positional encoding to a sequence
          position. Applies only if layer is created in ``'train'`` mode.
      dropout_broadcast_dims: Axes along which dropout mask values are
          broadcast rather than individually set at random.
      use_bfloat16: If ``True``, use bfloat16 weights instead of the default
        float32; this can save memory but may (rarely) lead to numerical issues.
      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).
      max_offset_to_add: maximum offset to add to the positions during training
        when randomizing; this offset plus input length must still be less than
        max_len for all training examples.
      d_feature: int or None; have this dimension for embeddings + shared FF if
        not None.
      mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
    """
    super().__init__()
    self._max_len = max_len
    if dropout >= 1.0:
      raise ValueError('Dropout rates must be lower than 1.')
    if mode == 'train':
      self._dropout = dropout
    else:
      self._dropout = 0.0
    self._dropout_broadcast_dims = dropout_broadcast_dims
    self._use_bfloat16 = use_bfloat16
    self._start_from_zero_prob = start_from_zero_prob
    self._max_offset_to_add = max_offset_to_add
    self._mode = mode
    self._d_feature = d_feature
Creates a :py:class:`PositionalEncoding` instance in a given mode.

    Args:
      max_len: Maximum input sequence length.
      dropout: Probability of *not* adding positional encoding to a sequence
          position. Applies only if layer is created in ``'train'`` mode.
      dropout_broadcast_dims: Axes along which dropout mask values are
          broadcast rather than individually set at random.
      use_bfloat16: If ``True``, use bfloat16 weights instead of the default
        float32; this can save memory but may (rarely) lead to numerical issues.
      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).
      max_offset_to_add: maximum offset to add to the positions during training
        when randomizing; this offset plus input length must still be less than
        max_len for all training examples.
      d_feature: int or None; have this dimension for embeddings + shared FF if
        not None.
      mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
    
trax/layers/attention.py