docs: def _fused_batch_norm_op(self, input_batch, mean, variance, use_batch_stats):
      """Creates a fused batch normalization op."""
      # Store the original shape of the mean and variance.
      mean_shape = mean.get_shape()
      variance_shape = variance.get_shape()
      # The fused batch norm expects the mean, variance, gamma and beta
      # tensors to have dimension 1, so we flatten them to remove the
      # extra dimensions. In addition, it expects the input_batch to have
      # dimension 4, so we reshape it accordingly.
      gamma_flatten = tf.reshape(self._gamma, shape=(self._num_channels,))
      beta_flatten = tf.reshape(self._beta, shape=(self._num_channels,))
      flatten_mean = tf.reshape(mean, shape=(self._num_channels,))
      flatten_variance = tf.reshape(variance, shape=(self._num_channels,))
      use_batch_stats = tf.convert_to_tensor(use_batch_stats)
  
      input_shape = input_batch.get_shape()
      output_shape = [-1] + input_shape.as_list()[1:]
  
      flat_image_size = np.prod(self._image_shape, dtype=np.int32)
  
      if len(self._data_format) == 4:
        fusable_data_format = self._data_format
        fusable_batch = input_batch
      elif self._channel_index == 1 and self._image_shape:
        fusable_data_format = "NCHW"
        fusable_batch = tf.reshape(
            input_batch,
            shape=(-1, self._num_channels, 1, flat_image_size))
      else:
        # The CPU implementation of FusedBatchNorm only supports NHWC tensor
        # format for now.
        fusable_data_format = "NHWC"
        fusable_batch = tf.reshape(
            input_batch,
            shape=(-1, 1, flat_image_size, self._num_channels))
  
      common_args = {
          "scale": gamma_flatten,
          "offset": beta_flatten,
          "epsilon": self._eps,
          "data_format": fusable_data_format,
          "name": "batch_norm"
      }
  
      def use_batch_stats_fused_batch_norm():
        return tf.nn.fused_batch_norm(
            fusable_batch,
            mean=None,
            variance=None,
            is_training=True,
            **common_args)
  
      def moving_average_fused_batch_norm():
        return tf.nn.fused_batch_norm(
            fusable_batch,
            mean=flatten_mean,
            variance=flatten_variance,
            is_training=False,
            **common_args)
  
      batch_norm_op, mean, variance = utils.smart_cond(
          use_batch_stats, use_batch_stats_fused_batch_norm,
          moving_average_fused_batch_norm)
  
      if len(self._data_format) != 4:
        batch_norm_op = tf.reshape(batch_norm_op, output_shape)
      mean = tf.reshape(mean, mean_shape)
      variance = tf.reshape(variance, variance_shape)
      return batch_norm_op, mean, variance