docs: def _build(self, input_batch, is_training, test_local_stats=True):
      """Connects the BatchNorm module into the graph.
  
      Args:
        input_batch: A Tensor of arbitrary dimension. By default, the final
          dimension is not reduced over when computing the minibatch statistics.
        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 `True`. Can be a Tensor.
  
      Returns:
        A tensor with the same shape as `input_batch`.
  
      Raises:
        base.IncompatibleShapeError: If `axis` is not valid for the
          input shape or has negative entries.
        base.NotSupportedError: If `input_batch` has data type of `tf.bfloat16`.
      """
      input_shape = input_batch.get_shape()
  
      if self._axis is not None:
        if len(self._axis) > len(input_shape):
          raise base.IncompatibleShapeError(
              "Too many indices specified in axis: len({}) > len({}).".format(
                  self._axis, input_shape))
  
        if max(self._axis) >= len(input_shape):
          raise base.IncompatibleShapeError(
              "One or more index in axis is too large for "
              "input shape: {} >= {:d}.".format(self._axis, len(input_shape)))
  
        if min(self._axis) < 0:
          raise base.IncompatibleShapeError(
              "Indices in axis must be non-negative: {} < 0.".format(
                  self._axis))
  
        axis = self._axis
      else:
        # Reduce over all dimensions except the last.
        axis = tuple(range(len(input_shape))[:-1])
  
      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._mean_shape = input_batch.get_shape().as_list()
      for index in axis:
        self._mean_shape[index] = 1
  
      use_batch_stats = is_training | test_local_stats
  
      mean, variance = self._build_statistics(input_batch, axis,
                                              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