code: 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=[],
      )
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import activation_fns
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
label_desc: Use image node_id 8 for calling a global function with example usage: SRU(n_units, activation, mode) and returns: cb