code: def _fused_batch_norm_op(self, input_batch, mean, variance, use_batch_stats):
      
      # 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.
      gamma_flatten = tf.reshape(self._gamma, shape=(-1,))
      beta_flatten = tf.reshape(self._beta, shape=(-1,))
      flatten_mean = tf.reshape(mean, shape=(-1,))
      flatten_variance = tf.reshape(variance, shape=(-1,))
      use_batch_stats = tf.convert_to_tensor(use_batch_stats)
  
      common_args = {
          "scale": gamma_flatten,
          "offset": beta_flatten,
          "epsilon": self._eps,
          "data_format": self._infer_fused_data_format(input_batch),
          "name": "batch_norm"
      }
  
      def use_batch_stats_fused_batch_norm():
        return tf.nn.fused_batch_norm(input_batch, mean=None, variance=None,
                                      is_training=True, **common_args)
  
      def moving_average_fused_batch_norm():
        return tf.nn.fused_batch_norm(input_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)
  
      mean = tf.reshape(mean, mean_shape)
      variance = tf.reshape(variance, variance_shape)
      return batch_norm_op, mean, variance
docs: Creates a fused batch normalization op.