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.
    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