Creates a batch normalization op.

    It uses the tf.nn.batch_normalization op by default and the
    tf.nn.fused_batch_norm op to support fused batch normalization.

    Args:
      input_batch: A input Tensor of arbitrary dimension.
      mean: A mean tensor, of the same dtype as `input_batch`.
      variance: A variance tensor, of the same dtype as `input_batch`.
      use_batch_stats: A bool value that indicates whether the operation should
         use the batch statistics.
      stat_dtype: TensorFlow datatype used for the moving mean and variance.

    Returns:
      A batch normalization operation.
      The current mean tensor, of datatype `stat_dtype`.
      The current variance tensor, of datatype `stat_dtype`.

def _batch_norm_op(self, input_batch, mean, variance, use_batch_stats,
                     stat_dtype):
    """Creates a batch normalization op.

    It uses the tf.nn.batch_normalization op by default and the
    tf.nn.fused_batch_norm op to support fused batch normalization.

    Args:
      input_batch: A input Tensor of arbitrary dimension.
      mean: A mean tensor, of the same dtype as `input_batch`.
      variance: A variance tensor, of the same dtype as `input_batch`.
      use_batch_stats: A bool value that indicates whether the operation should
         use the batch statistics.
      stat_dtype: TensorFlow datatype used for the moving mean and variance.

    Returns:
      A batch normalization operation.
      The current mean tensor, of datatype `stat_dtype`.
      The current variance tensor, of datatype `stat_dtype`.
    """
    if self._fused:
      # For the non-training case where not using batch stats,
      # pass in the moving statistic variables directly.
      # These will already be in the correct dtype, even for float16 input.
      batch_norm_op, mean, variance = self._fused_batch_norm_op(
          input_batch,
          self._moving_mean, self._moving_variance, use_batch_stats)
    else:
      batch_norm_op = tf.nn.batch_normalization(
          input_batch,
          mean,
          variance,
          self._beta,
          self._gamma,
          self._eps,
          name="batch_norm")
      # We'll echo the supplied mean and variance so that they can also be used
      # to update the moving statistics. Cast to matching type if necessary.
      if input_batch.dtype.base_dtype != stat_dtype:
        mean = tf.cast(mean, stat_dtype)
        variance = tf.cast(variance, stat_dtype)

    return batch_norm_op, mean, variance