code: 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,
          ),
      )
  
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 7 for calling a global function with example usage: ScanSRUCell(mode, monkey_patched_mask) and returns: cb, cb, updated_mask,  x_times_one_minus_f, res,  res