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