def init_weights_and_state(self, input_signature):
    """Initializes weights and state for inputs with the given signature."""
    states = []
    for sublayer in [self._lsha, self._stda]:
      unused_weights_or_cache_marker, state_or_cache_marker = sublayer.init(
          input_signature, use_cache=False)
      states.append(state_or_cache_marker)
    self.state = tuple(states)
    self.weights = ()  # Wrapper forward_and_backward assumes this is ()
Initializes weights and state for inputs with the given signature.
trax/layers/research/efficient_attention.py