code: def ConvGRUCell(n_units, kernel_size=(3, 3)):
      """Builds a convolutional GRU.
  
      Paper: https://arxiv.org/abs/1511.06432.
  
      Args:
        n_units: Number of hidden units
        kernel_size: Kernel size for convolution
  
      Returns:
        A Stax model representing a GRU cell with convolution transforms.
      """
  
      def BuildConv():
          return convolution.Conv(
              filters=n_units, kernel_size=kernel_size, padding="SAME"
          )
  
      return GeneralGRUCell(
          candidate_transform=BuildConv,
          memory_transform_fn=None,
          gate_nonlinearity=activation_fns.Sigmoid,
          candidate_nonlinearity=activation_fns.Tanh,
      )
  
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 4 for calling a global function with example usage: ConvGRUCell(n_units, kernel_size) and returns: GeneralGRUCell, convolution