code: def GRU(n_units, mode="train"):
      """GRU running on axis 1."""
      zero_state = MakeZeroState(
          depth_multiplier=1
      )  # pylint: disable=no-value-for-parameter
      return cb.Serial(
          cb.Branch([], zero_state),
          cb.Scan(GRUCell(n_units=n_units), axis=1, mode=mode),
          cb.Select([0], n_in=2),  # Drop RNN state.
          # Set the name to GRU and don't print sublayers.
          name=f"GRU_{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 3 for calling a global function with example usage: GRU(n_units, mode) and returns: cb