code: def ConvGRUCell(units, kernel_size=(3, 3)):
    
  
    def BuildConv():
      return core.Conv(filters=units, kernel_size=kernel_size, padding='SAME')
  
    return GeneralGRUCell(
        candidate_transform=BuildConv,
        memory_transform=combinators.Identity,
        gate_nonlinearity=core.Sigmoid,
        candidate_nonlinearity=core.Tanh)
docs: Builds a convolutional GRU.
  
    Paper: https://arxiv.org/abs/1511.06432.
  
    Args:
      units: Number of hidden units
      kernel_size: Kernel size for convolution
  
    Returns:
      A Stax model representing a GRU cell with convolution transforms.