def init_weights_and_state(self, input_signature):
    """Randomly initializes the positional encoding vectors.

    Args:
      input_signature: `ShapeDtype` instance characterizing the input this
          layer should compute on.
    """
    if self._mode == 'predict':
      self.state = jnp.zeros((), dtype=jnp.int32)
Randomly initializes the positional encoding vectors.

    Args:
      input_signature: `ShapeDtype` instance characterizing the input this
          layer should compute on.
    
trax/layers/research/position_encodings.py