def SRU(n_units, activation=None, mode='train'):
  r"""SRU (Simple Recurrent Unit) layer as in https://arxiv.org/abs/1709.02755.

  As defined in the paper:

  .. math::
    y_t &= W x_t + B \quad \hbox{(include $B$ optionally)} \\
    f_t &= \sigma(Wf x_t + bf) \\
    r_t &= \sigma(Wr x_t + br) \\
    c_t &= f_t \times c_{t-1} + (1 - f_t) \times y_t \\
    h_t &= r_t \times \hbox{activation}(c_t) + (1 - r_t) \times x_t

  We assume the input is of shape [batch, length, depth] and recurrence
  happens on the length dimension. This returns a single layer. It's best
  to use at least 2, they say in the paper, except inside a Transformer.

  Args:
    n_units: output depth of the SRU layer.
    activation: Optional activation function.
    mode: if 'predict' then we save the previous state for one-by-one inference

  Returns:
    The SRU layer.
  """
  sigmoid_activation = activation_fns.Sigmoid()
  return cb.Serial(                                         # x
      cb.Branch(core.Dense(3 * n_units), []),               # r_f_y, x
      cb.Split(n_items=3),                                  # r, f, y, x
      cb.Parallel(sigmoid_activation, sigmoid_activation),  # r, f, y, x
      base.Fn('',
              lambda r, f, y: (y * (1.0 - f), f, r),    # y * (1 - f), f, r, x
              n_out=3),
      cb.Parallel([], [], cb.Branch(MakeZeroState(), [])),
      ScanSRUCell(mode=mode),
      cb.Select([0], n_in=2),                               # act(c), r, x
      activation if activation is not None else [],
      base.Fn('FinalSRUGate', lambda c, r, x: c * r + x * (1 - r) * (3**0.5)),
      # Set the name to SRU and don't print sublayers.
      name=f'SRU_{n_units}', sublayers_to_print=[]
  )
SRU (Simple Recurrent Unit) layer as in https://arxiv.org/abs/1709.02755.

  As defined in the paper:

  .. math::
    y_t &= W x_t + B \quad \hbox{(include $B$ optionally)} \\
    f_t &= \sigma(Wf x_t + bf) \\
    r_t &= \sigma(Wr x_t + br) \\
    c_t &= f_t \times c_{t-1} + (1 - f_t) \times y_t \\
    h_t &= r_t \times \hbox{activation}(c_t) + (1 - r_t) \times x_t

  We assume the input is of shape [batch, length, depth] and recurrence
  happens on the length dimension. This returns a single layer. It's best
  to use at least 2, they say in the paper, except inside a Transformer.

  Args:
    n_units: output depth of the SRU layer.
    activation: Optional activation function.
    mode: if 'predict' then we save the previous state for one-by-one inference

  Returns:
    The SRU layer.
  
trax/layers/rnn.py