def __init__(self, time_bin_length, mode='train'):
    """Initializes the encoding.

    Args:
      time_bin_length: TimeBinCausalAttention.bin_length of the first layer.
      mode: if 'predict', allow evaluating one token at a time
    """
    super().__init__()
    self._time_bin_length = time_bin_length
    self._mode = mode
Initializes the encoding.

    Args:
      time_bin_length: TimeBinCausalAttention.bin_length of the first layer.
      mode: if 'predict', allow evaluating one token at a time
    
trax/layers/research/position_encodings.py