def __init__(self, add_offset=2048, dropout=0.0, dropout_broadcast_dims=(-2,),
               start_from_zero_one_in=2, mode='train'):
    """Creates a SinCosPositionalEncoding instance.

    Args:
      add_offset: Maximumnumber to add to positions during training.
      dropout: Probability of *not* adding positional encoding to a sequence
          position.
      dropout_broadcast_dims: Axes along which dropout mask values are
          broadcast rather than individually set at random.
      start_from_zero_one_in: how often to start from 0 during training,
          every one in that many times (e.g., if 4, then it's 25% of the time).
      mode: One of `'train'`, `'eval'`, or `'predict'`.
    """
    super().__init__()
    self._add_offset = add_offset
    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._start_from_zero_one_in = start_from_zero_one_in
    self._mode = mode
Creates a SinCosPositionalEncoding instance.

    Args:
      add_offset: Maximumnumber to add to positions during training.
      dropout: Probability of *not* adding positional encoding to a sequence
          position.
      dropout_broadcast_dims: Axes along which dropout mask values are
          broadcast rather than individually set at random.
      start_from_zero_one_in: how often to start from 0 during training,
          every one in that many times (e.g., if 4, then it's 25% of the time).
      mode: One of `'train'`, `'eval'`, or `'predict'`.
    
trax/layers/research/position_encodings.py