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 = _make_tuple(input_signature)

    # Inputs/outputs of `cond`.
    inputs = inputs_from_stack(stack, self._cond.n_in)
    weights_or_cache_marker, state_or_cache_marker = (
        self._cond.init(inputs, use_cache=True))
    weights.append(weights_or_cache_marker)
    states.append(state_or_cache_marker)
    self._cond._forward_abstract(inputs)
    stack = _make_tuple(outputs_onto_stack([], stack, self._cond.n_in))

    # Inputs/outputs of `true` and `false`.
    for sublayer in [self._true, self._false]:
      inputs = inputs_from_stack(stack, sublayer.n_in)
      weights_or_cache_marker, state_or_cache_marker = (
          sublayer.init(inputs, use_cache=True))
      weights.append(weights_or_cache_marker)
      states.append(state_or_cache_marker)

    self.state = states
    self.weights = weights
    # pylint: enable=protected-access
Initializes weights and state for inputs with the given signature.
trax/layers/combinators.py