def forward(self, inputs):
    """Executes this layer as part of a forward pass through the model."""
    weights = self.weights[0]
    if isinstance(inputs, list):
      inputs = tuple(inputs)  # so that inputs structure matches outputs
    n_carry = self._n_carry
    def scannable_fn(x, carry_and_state):  # pylint: disable=invalid-name
      carry, state, i = carry_and_state
      x_and_carry = x + carry if n_carry > 0 else x
      rng = fastmath.random.fold_in(self.rng, i)
      res, new_state = self.sublayer.pure_fn(
          x_and_carry, weights, state, rng, use_cache=True)
      if n_carry > 0:
        return (res[:-n_carry], (res[-n_carry:], new_state, i+1))
      else:
        return (res, ([], new_state, i+1))

    if n_carry > 0:
      xs = inputs[:-n_carry]  # Split input stack into inputs and carry.
      xs_carry = inputs[-n_carry:]
      if self._mode == 'predict' and self._state[1] is not ():  # pylint: disable=literal-comparison
        xs_carry = self._state[1]
      init = (xs_carry, self.state[0], jnp.array(0, dtype=jnp.int32))
    else:
      xs_carry = ()
      xs, init = inputs, ([], self.state[0], jnp.array(0, dtype=jnp.int32))
    ys, (carry, new_state, _) = _scan(scannable_fn, xs, init,
                                      axis=self._axis, remat=self._remat)
    res = ys + carry if n_carry > 0 else ys
    state_carry = carry if self._mode == 'predict' and n_carry > 0 else ()
    self.state = (new_state, state_carry)
    return res  # Put outputs and carry back on stack.
Executes this layer as part of a forward pass through the model.
trax/layers/combinators.py