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.

def ConvGRUCell(units, kernel_size=(3, 3)):
  """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.
  """

  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)