def init_weights_and_state(self, input_signature):
    """Initializes weights and state for inputs with the given signature."""
    if self._n_in == 1 and not isinstance(input_signature, (list, tuple)):
      input_signature = (input_signature,)
    batched_signature = []
    for sub_input_signature in input_signature:
      batched_size = 1
      for d in sub_input_signature.shape[:-self._n_last_axes_to_keep]:
        batched_size *= d
      batched_shape = [batched_size] + list(
          sub_input_signature.shape[-self._n_last_axes_to_keep:])
      batched_signature.append(ShapeDtype(batched_shape,
                                          sub_input_signature.dtype))
    if self._n_in == 1:
      batched_signature = batched_signature[0]
    weights, layer_state = self.sublayer.init(batched_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