Connects the BatchNormV2 module into the graph.

    Args:
      input_batch: A Tensor of the same dimension as `len(data_format)`.
      is_training: A boolean to indicate if the module should be connected in
        training mode, meaning the moving averages are updated. Can be a Tensor.
      test_local_stats: A boolean to indicate if local batch statistics should
        be used when `is_training=False`. If not, moving averages are used.
        By default `False`. Can be a Tensor.

    Returns:
      A tensor with the same shape as `input_batch`.

    Raises:
      base.IncompatibleShapeError: If `data_format` is not valid for the
        input shape.
      base.NotSupportedError: If `input_batch` has data type of `tf.bfloat16`.

def _build(self,
             input_batch,
             is_training,
             test_local_stats=False):
    """Connects the BatchNormV2 module into the graph.

    Args:
      input_batch: A Tensor of the same dimension as `len(data_format)`.
      is_training: A boolean to indicate if the module should be connected in
        training mode, meaning the moving averages are updated. Can be a Tensor.
      test_local_stats: A boolean to indicate if local batch statistics should
        be used when `is_training=False`. If not, moving averages are used.
        By default `False`. Can be a Tensor.

    Returns:
      A tensor with the same shape as `input_batch`.

    Raises:
      base.IncompatibleShapeError: If `data_format` is not valid for the
        input shape.
      base.NotSupportedError: If `input_batch` has data type of `tf.bfloat16`.
    """
    input_shape = input_batch.get_shape()

    if not self._data_format:
      if len(input_shape) == 2:
        self._data_format = "NC"
      elif len(input_shape) == 3:
        self._data_format = "NWC"
      elif len(input_shape) == 4:
        self._data_format = "NHWC"
      elif len(input_shape) == 5:
        self._data_format = "NDHWC"
      else:
        raise base.IncompatibleShapeError(
            "Input shape {} has too many or too few dimensions.".format(
                input_shape))

    self._channel_index = self._data_format.index("C")
    # Use list to turn range into iterator in python3.
    self._axis = list(range(len(self._data_format)))
    del self._axis[self._channel_index]

    if len(self._data_format) != len(input_shape):
      raise base.IncompatibleShapeError(
          "Incorrect data format {} for input shape {}.".format(
              self._data_format, input_shape))

    dtype = input_batch.dtype.base_dtype
    if self._fused and dtype == tf.bfloat16:
      raise base.NotSupportedError(
          "Fused batch norm does not support tf.bfloat16.")
    # Maintain moving averages at a minimum precision of tf.float32.
    stat_dtype = tf.float32 if dtype in [tf.float16, tf.bfloat16] else dtype

    self._num_channels = int(input_shape[self._channel_index])
    if self._channel_index == 1:
      self._image_shape = [int(x) for x in input_shape[2:]]
    else:
      self._image_shape = [int(x) for x in input_shape[1:-1]]

    self._expanded_mean_shape = [1] * len(input_shape)
    self._expanded_mean_shape[self._channel_index] = self._num_channels

    use_batch_stats = is_training | test_local_stats

    mean, variance = self._build_statistics(input_batch, use_batch_stats,
                                            stat_dtype)

    # Sets up optional gamma and beta parameters
    self._build_scale_offset(dtype)
    # Sets up the batch normalization op.
    out, mean, variance = self._batch_norm_op(input_batch, mean, variance,
                                              use_batch_stats, stat_dtype)
    # Sets up the update op.
    update_ops = self._build_update_ops(mean, variance, is_training)

    # Put update ops in the update ops collection if given, otherwise add as
    # control dependencies of the output.
    if update_ops:
      if self._update_ops_collection:
        for update_op in update_ops:
          tf.add_to_collection(self._update_ops_collection, update_op)
      else:
        with tf.control_dependencies(update_ops):
          out = tf.identity(out)

    return out