def init_weights_and_state(self, input_signature):
    """Initializes weights and state for inputs with the given signature."""
    weights, layer_state = self.sublayer.init(input_signature, use_cache=True)
    self.state = ((), layer_state)
    self._weights = (weights,)
Initializes weights and state for inputs with the given signature.
trax/layers/combinators.py