code: def ConvDiagonalGRU(units, kernel_size=(3, 3)):
    
  
    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)
docs: Build convolutional GRU with diagonal gating as in ImprovedNGPU.