def ConvDiagonalGRU(units, kernel_size=(3, 3)):
  """Build convolutional GRU with diagonal gating as in ImprovedNGPU."""

  def BuildConv():
    return layers.Conv(filters=units, kernel_size=kernel_size, padding='SAME')

  return layers.GeneralGRUCell(
      candidate_transform=BuildConv,
      memory_transform=DiagonalGate,
      gate_nonlinearity=layers.HardSigmoid,
      candidate_nonlinearity=layers.HardTanh)