docs: 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