WideResnet convolutational block.

def WideResnetBlock(channels, strides=(1, 1), channel_mismatch=False):
  """WideResnet convolutational block."""
  main = layers.Serial(layers.BatchNorm(), layers.Relu(),
                       layers.Conv(channels, (3, 3), strides, padding='SAME'),
                       layers.BatchNorm(), layers.Relu(),
                       layers.Conv(channels, (3, 3), padding='SAME'))
  shortcut = layers.Identity() if not channel_mismatch else layers.Conv(
      channels, (3, 3), strides, padding='SAME')
  return layers.Serial(
      layers.Branch(), layers.Parallel(main, shortcut), layers.SumBranches())