def forward(self, xs):
    """Executes this layer as part of a forward pass through the model.

    Args:
      xs: Tensors of as required by the branches of this conditional.

    Returns:
      Tensors resulting from running the chosen branch.
    """
    # TODO(jaszczur): modify; it's a copy from SkippingSerial
    self._validate_forward_inputs(xs)
    layers_state = self.state
    # Get 3 rngs, one for each layer.
    rngs = _split_rngs(self.rng, 3)

    # Prepare the stack and do some safety checks as in the parent class.
    stack = _make_tuple(xs)
    weights = self.weights
    if len(weights) != 3:
      raise ValueError('number of weights ({}) not equal to 3'
                       .format(len(weights)))
    if len(layers_state) != 3:
      raise ValueError('length of state ({}) not equal to 3'
                       .format(len(layers_state)))

    def true_func(t):
      outputs, new_true_state = self._true.pure_fn(
          t[0][0], t[1][0], t[2][0], t[3][0])
      # t[2][1] is old_false_state which is not changing if true is executed.
      return outputs, (new_true_state, t[2][1])

    def false_func(t):
      if self._identity_false_fun:
        # Memory optimization: we don't need pure_fn call.
        return t[0][1], t[2]
      outputs, new_false_state = self._false.pure_fn(
          t[0][1], t[1][1], t[2][1], t[3][1])
      # t[2][1] is old_true_state, which is not changing if false is executed.
      return outputs, (t[2][0], new_false_state)

    cond_inputs = inputs_from_stack(xs, self._cond.n_in)
    cond_output, s = self._cond.pure_fn(cond_inputs, self.weights[0],
                                        self.state[0], rngs[0], use_cache=True)
    stack = outputs_onto_stack([], stack, self._cond.n_in)
    self._cond.state = s

    outputs, both_states = fastmath.cond(
        cond_output,
        true_func,
        false_func,
        [(stack, stack),
         (self.weights[1], self.weights[2]),
         (self.state[1], self.state[2]),
         (rngs[1], rngs[2])]
    )
    stack = outputs_onto_stack([], stack, self._cond.n_in)

    # We don't know which (`true` or `false`) branch was run, but both of them
    # are adding (n_out) and removing (n_in) the same number of elements of the
    # stack (this was checked in __init__). outputs_onto_stack just uses the
    # layer's n_in, so we can pass either `true` or `false` to it.
    # Note that `outputs` is the actual output of `true` or `false` branch,
    # whichever was run, and we add it to the stack in any case.
    stack = outputs_onto_stack(outputs, stack, self._true.n_in)
    self._true.state = both_states[0]
    self._false.state = both_states[1]
    return _make_singleitem_or_original(stack)
Executes this layer as part of a forward pass through the model.

    Args:
      xs: Tensors of as required by the branches of this conditional.

    Returns:
      Tensors resulting from running the chosen branch.
    
trax/layers/combinators.py