def init_weights_and_state(self, input_signature):
    """Initializes weights and state for inputs with the given signature."""
    weights = []
    states = []
    # In the code below, stack, inputs, and outputs are abstract (shapes and
    # dtypes), but weights and states are non-abstract actual values.
    stack = input_signature
    for sublayer in self.sublayers:
      inputs = inputs_from_stack(stack, sublayer.n_in)
      weights_or_cache_marker, state_or_cache_marker = (
          sublayer.init(inputs, use_cache=True))
      outputs, _ = sublayer._forward_abstract(inputs)
      stack = outputs_onto_stack(outputs, stack, sublayer.n_in)

      weights.append(weights_or_cache_marker)
      states.append(state_or_cache_marker)
    self.state = tuple(states)
    self.weights = tuple(weights)
Initializes weights and state for inputs with the given signature.
trax/layers/combinators.py