def init_weights_and_state(self, input_signature):
    """Initializes this layer for fast inference, if in ``'predict'`` mode."""
    if self._mode == 'predict':
      self.state = _fast_inference_init_state(
          input_signature, self._max_len,
          predict_mask=self._portal_mask)
Initializes this layer for fast inference, if in ``'predict'`` mode.
trax/layers/attention.py