Creates a fused batch normalization op.

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