def __init__(
      self, drift=.03, affine=True, transform='any',
      time_bin_length=None,
      mode='train'):
    """Initializes the encoding.

    The encoding tries to roughly evenly traverse the latent space.
    The recurrence time is dependent on how many bits per dimension you use.

    There are two parameters to control randomization:
    - randomizing the origin every 1/drift steps by letting it drift
    - randomizing the origin per call

    Args:
      drift: variance in position difference per unit of difference
      affine: whether to randomize the origin every call
      transform: learnable transform after encoding (any/diag/none)
      time_bin_length: Add features AxialPositionalEncoding learns if
        TimeBinCausalAttention is the first layer.
        bin_length should match TBCA.bin_length
        If you set transform='diag', this flag increases your model capacity to
        close to transform='any', though it will still train slower.
      mode: if 'predict', allow evaluating one token at a time
    """
    super().__init__()
    if transform not in ('any', 'diag', 'none'):
      raise ValueError(transform)
    self._noise_rng = jax.random.split(jax.random.PRNGKey(234234535))[0]
    assert self._noise_rng is not None
    self._noise = None
    self._drift = drift
    self._affine = affine
    self._transform = transform
    self._time_bin_length = time_bin_length
    self._mode = mode
Initializes the encoding.

    The encoding tries to roughly evenly traverse the latent space.
    The recurrence time is dependent on how many bits per dimension you use.

    There are two parameters to control randomization:
    - randomizing the origin every 1/drift steps by letting it drift
    - randomizing the origin per call

    Args:
      drift: variance in position difference per unit of difference
      affine: whether to randomize the origin every call
      transform: learnable transform after encoding (any/diag/none)
      time_bin_length: Add features AxialPositionalEncoding learns if
        TimeBinCausalAttention is the first layer.
        bin_length should match TBCA.bin_length
        If you set transform='diag', this flag increases your model capacity to
        close to transform='any', though it will still train slower.
      mode: if 'predict', allow evaluating one token at a time
    
trax/layers/research/position_encodings.py