def init_weights_and_state(self, input_signature):
    """Initializes weights and state for inputs with the given signature."""
    sublayer_signatures = self._allot_to_sublayers(input_signature)
    inits = [layer.init(signature, use_cache=True)
             for layer, signature
             in zip(self.sublayers, sublayer_signatures)]
    if inits:
      weights, state = tuple(zip(*inits))
      self.state = state
      self.weights = weights
Initializes weights and state for inputs with the given signature.
trax/layers/combinators.py