def forward(self, inputs):
    """Executes this layer as part of a forward pass through the model."""
    n_layers, layers = self._n_layers, self.sublayers
    sublayer_inputs = self._allot_to_sublayers(inputs)
    state, weights = self.state, self.weights
    rngs = _split_rngs(self.rng, n_layers)
    if len(sublayer_inputs) != n_layers:
      raise ValueError(
          f'Number of inputs for sublayers ({len(sublayer_inputs)}) does not equal '
          f'number of sublayers ({n_layers}).')
    if len(weights) != n_layers:
      raise ValueError(
          f'Number of weight elements ({len(weights)}) does not equal '
          f'number of sublayers ({n_layers}).')
    if len(state) != n_layers:
      raise ValueError(
          f'Number of state elements ({len(state)}) does not equal '
          f'number of sublayers ({n_layers}).')
    if len(rngs) != n_layers:
      raise ValueError(
          f'Number of rngs ({len(rngs)}) does not equal '
          f'number of sublayers ({n_layers}).')
    outputs = []
    new_state = []
    for layer, x, w, s, r in zip(layers, sublayer_inputs, weights, state, rngs):
      # Note that zip silently truncates its result if lengths don't match.
      sub_outputs, sub_state = layer.pure_fn(x, w, s, r, use_cache=True)
      if layer.n_out == 1:
        outputs.append(sub_outputs)
      else:
        outputs.extend(sub_outputs)
      new_state.append(sub_state)
    output = outputs[0] if self.n_out == 1 else tuple(outputs)
    self.state = tuple(new_state)
    return output
Executes this layer as part of a forward pass through the model.
trax/layers/combinators.py