def forward(self, xs):
    """Executes this layer as part of a forward pass through the model."""
    self._validate_forward_inputs(xs)
    if not self.sublayers:  # No-op: outputs = inputs
      return xs

    state, weights = self.state, self.weights
    rngs = _split_rngs(self.rng, self._n_layers)
    stack = xs
    new_state = []
    n_layers = self._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}).')

    for layer, w, s, rng in zip(self.sublayers, weights, state, rngs):
      inputs = inputs_from_stack(stack, layer.n_in)
      outputs, s = layer.pure_fn(inputs, w, s, rng, use_cache=True)
      stack = outputs_onto_stack(outputs, stack, layer.n_in)
      new_state.append(s)
    self.state = tuple(new_state)
    return stack
Executes this layer as part of a forward pass through the model.
trax/layers/combinators.py