def InnerSRUCell():
  """The inner (non-parallel) computation of an SRU."""
  def f(cur_x_times_one_minus_f, cur_f, cur_state):  # pylint: disable=invalid-name
    res = cur_f * cur_state + cur_x_times_one_minus_f
    return res, res
  return base.Fn('InnerSRUCell', f, n_out=2)
The inner (non-parallel) computation of an SRU.
trax/layers/rnn.py