def init_weights_and_state(self, input_signature):
    """Initializes weights and state for inputs with the given signature."""
    n_carry = self._n_carry
    if n_carry == 0:
      if isinstance(input_signature, (list, tuple)):
        layer_sig = [ShapeDtype(_shape_without_axis(x, self._axis), x.dtype)
                     for x in input_signature]
        layer_sig = tuple(layer_sig)
      else:
        layer_sig = ShapeDtype(_shape_without_axis(input_signature, self._axis),
                               input_signature.dtype)
      weights, state = self.sublayer.init(layer_sig)
      self.state = (state, ())
      self.weights = (weights,)
    else:
      xs = input_signature[:-n_carry]
      init = input_signature[-n_carry:]
      xs_slices = [ShapeDtype(_shape_without_axis(x, self._axis), x.dtype)
                   for x in xs]
      layer_signature = tuple(xs_slices + list(init))
      weights, state = self.sublayer.init(layer_signature, use_cache=True)
      self.state = (state, ())
      self.weights = (weights,)
Initializes weights and state for inputs with the given signature.
trax/layers/combinators.py