def ScanSRUCell(mode, monkey_patched_mask=None):
  """The inner (non-parallel) computation of an SRU."""
  if monkey_patched_mask is None:
    return cb.Scan(InnerSRUCell(), axis=1, mode=mode)

  # This is necessary for Terraformer model. See comments there.
  # The mask will only be used in Terraformer in predict mode.
  assert mode == 'predict'

  def update_mask(mask, x_times_one_minus_f):  # pylint: disable=invalid-name
    initial = jnp.ones(x_times_one_minus_f.shape[:2], dtype=jnp.float32)
    if initial.shape[1] > 1:
      updated_mask = fastmath.dynamic_update_slice_in_dim(
          initial != 0, mask != 0, 1, axis=1)
    else:
      updated_mask = initial
    return updated_mask, x_times_one_minus_f

  def masked_inner_sru_cell(cur_mask, cur_x_times_one_minus_f, cur_f,  # pylint: disable=invalid-name
                            cur_state):
    res = ((cur_f * cur_state + cur_x_times_one_minus_f) * cur_mask
           + (1 - cur_mask) * cur_state)
    return res, res

  return cb.Serial(
      monkey_patched_mask.get_layer(),
      base.Fn('update_mask', update_mask, n_out=2),
      cb.Scan(base.Fn('MaskedInnerSRUCell', masked_inner_sru_cell, n_out=2),
              axis=1, mode=mode),
      )
The inner (non-parallel) computation of an SRU.
trax/layers/rnn.py