def WideResnet(n_blocks=3, widen_factor=1, n_output_classes=10, bn_momentum=0.9,
               mode='train'):
  """WideResnet from https://arxiv.org/pdf/1605.07146.pdf.

  Args:
    n_blocks: int, number of blocks in a group. total layers = 6n + 4.
    widen_factor: int, widening factor of each group. k=1 is vanilla resnet.
    n_output_classes: int, number of distinct output classes.
    bn_momentum: float, momentum in BatchNorm.
    mode: Whether we are training or evaluating or doing inference.

  Returns:
    The list of layers comprising a WideResnet model with the given parameters.
  """
  return tl.Serial(
      tl.ToFloat(),
      tl.Conv(16, (3, 3), padding='SAME'),
      WideResnetGroup(n_blocks, 16 * widen_factor, bn_momentum=bn_momentum,
                      mode=mode),
      WideResnetGroup(n_blocks, 32 * widen_factor, (2, 2),
                      bn_momentum=bn_momentum, mode=mode),
      WideResnetGroup(n_blocks, 64 * widen_factor, (2, 2),
                      bn_momentum=bn_momentum, mode=mode),
      tl.BatchNorm(momentum=bn_momentum, mode=mode),
      tl.Relu(),
      tl.AvgPool(pool_size=(8, 8)),
      tl.Flatten(),
      tl.Dense(n_output_classes),
  )
WideResnet from https://arxiv.org/pdf/1605.07146.pdf.

  Args:
    n_blocks: int, number of blocks in a group. total layers = 6n + 4.
    widen_factor: int, widening factor of each group. k=1 is vanilla resnet.
    n_output_classes: int, number of distinct output classes.
    bn_momentum: float, momentum in BatchNorm.
    mode: Whether we are training or evaluating or doing inference.

  Returns:
    The list of layers comprising a WideResnet model with the given parameters.
  
trax/models/resnet.py