Python tensorflow.unique_with_counts() Examples

def get_distances(features, labels, num_classes):
    len_features = features.get_shape()[1]
    centers = tf.get_variable('centers', [num_classes, len_features], dtype=tf.float32,
                              initializer=tf.constant_initializer(0), trainable=False)
    labels = tf.reshape(labels, [-1])
    centers_batch = tf.gather(centers, labels)

    # distances = features - centers_batch
    diff = centers_batch - features
    unique_label, unique_idx, unique_count = tf.unique_with_counts(labels)
    appear_times = tf.gather(unique_count, unique_idx)
    appear_times = tf.reshape(appear_times, [-1, 1])

    diff = tf.divide(diff, tf.cast((1 + appear_times), tf.float32))

    return diff 

def repeat_with_index(x: tf.Tensor, index: tf.Tensor, axis: int = 1):
    """
    Given an tensor x (N*M*K), repeat the middle axis (axis=1)
    according to the index tensor index (G, )
    for example, if axis=1 and n = Tensor([0, 0, 0, 1, 2, 2])
    then M = 3 (3 unique values),
    and the final tensor would have the shape (N*6*3) with the
    first one in M repeated 3 times,
    second 1 time and third 2 times.

     Args:
        x: (3d Tensor) tensor to be augmented
        index: (1d Tensor) repetition tensor
        axis: (int) axis for repetition
    Returns:
        (3d Tensor) tensor after repetition
    """
    index = tf.reshape(index, (-1,))
    _, _, n = tf.unique_with_counts(index)
    return _repeat(x, n, axis) 

def knn(X_test, X_ref, Y_ref, K = 5):
	

	nearest_neighbors=tf.Variable(tf.zeros([K]))

	distance = tf.negative(tf.reduce_sum(tf.abs(tf.subtract(X_ref, X_test[0])),axis=1)) #L1
	values,indices=tf.nn.top_k(distance,k=K,sorted=False)

	nn = []
	
	for k in range(K):
		nn.append(tf.argmax(Y_ref[indices[k]], 0)) 

	nearest_neighbors=nn
	y, idx, count = tf.unique_with_counts(nearest_neighbors)

	preds = tf.slice(y, begin=[tf.argmax(count, 0)], size=tf.constant([1], dtype=tf.int64))[0]
		
	return preds 

def sparse_tensor_left_align(sparse_tensor):
  """Re-arranges a `tf.SparseTensor` and returns a left-aligned version of it.

  This mapper can be useful when returning a sparse tensor that may not be
  left-aligned from a preprocessing_fn.

  Args:
    sparse_tensor: A `tf.SparseTensor`.

  Returns:
    A left-aligned version of sparse_tensor as a `tf.SparseTensor`.
  """
  reordered_tensor = tf.sparse.reorder(sparse_tensor)
  transposed_indices = tf.transpose(reordered_tensor.indices)
  row_indices = transposed_indices[0]
  row_counts = tf.unique_with_counts(row_indices, out_idx=tf.int64).count
  column_indices = tf.ragged.range(row_counts).flat_values
  return tf.SparseTensor(
      indices=tf.transpose(tf.stack([row_indices, column_indices])),
      values=reordered_tensor.values,
      dense_shape=reordered_tensor.dense_shape) 

def knn(X_test, X_ref, Y_ref, K = 5):
	

	nearest_neighbors=tf.Variable(tf.zeros([K]))

	distance = tf.negative(tf.reduce_sum(tf.abs(tf.subtract(X_ref, X_test[0])),axis=1)) #L1
	values,indices=tf.nn.top_k(distance,k=K,sorted=False)

	nn = []
	
	for k in range(K):
		nn.append(tf.argmax(Y_ref[indices[k]], 0)) 

	nearest_neighbors=nn
	y, idx, count = tf.unique_with_counts(nearest_neighbors)

	preds = tf.slice(y, begin=[tf.argmax(count, 0)], size=tf.constant([1], dtype=tf.int64))[0]
		
	return preds 

def reduce_batch_count_or_sum_per_key(x, key, reduce_instance_dims):
  """Computes per-key sums or counts in the given tensor.

  Args:
    x: A `Tensor` or `SparseTensor`.  If x is None, return count per key.
    key: A `Tensor` or `SparseTensor` (cannot be None).
        Must meet one of the following conditions:
        1. Both x and key are dense,
        2. Both x and key are sparse and `key` must exactly match `x` in
        everything except values,
        3. The axis=1 index of each x matches its index of dense key.
    reduce_instance_dims: A bool, if True - collapses the batch and instance
        dimensions to arrive at a single scalar output. Otherwise, only
        collapses the batch dimension and outputs a `Tensor` of the same shape
        as the input. Not supported for `SparseTensor`s.

  Returns:
    A 2-tuple containing the `Tensor`s (key_vocab, count-or-sum).
  """
  if isinstance(x, tf.SparseTensor) and not reduce_instance_dims:
    raise NotImplementedError(
        'Sum per key only supports reduced dims for SparseTensors')

  key = _to_string(key)

  if x is not None:
    x, key = _validate_and_get_dense_value_key_inputs(x, key)
    unique = tf.unique(key, out_idx=tf.int64)
    if reduce_instance_dims and x.get_shape().ndims > 1:
      sums = tf.math.reduce_sum(x, axis=list(range(1, x.get_shape().ndims)))
    else:
      sums = x
    sums = tf.math.unsorted_segment_sum(sums, unique.idx, tf.size(unique.y))
  else:
    if isinstance(key, tf.SparseTensor):
      key = key.values
    key.set_shape([None])
    unique = tf.unique_with_counts(key, out_idx=tf.int64)
    sums = unique.count

  return unique.y, sums 

def get_git_loss(features, labels, num_classes):
    len_features = features.get_shape()[1]
    centers = tf.get_variable('centers', [num_classes, len_features], dtype=tf.float32,
                              initializer=tf.constant_initializer(0), trainable=False)
    labels = tf.reshape(labels, [-1])
    centers_batch = tf.gather(centers, labels)

    loss = tf.reduce_mean(tf.square(features - centers_batch))

    # Pairwise differences
    diffs = (features[:, tf.newaxis] - centers_batch[tf.newaxis, :])
    diffs_shape = tf.shape(diffs)

    # Mask diagonal (where i == j)
    mask = 1 - tf.eye(diffs_shape[0], diffs_shape[1], dtype=diffs.dtype)
    diffs = diffs * mask[:, :, tf.newaxis]

    # combinaton of two losses
    loss2 = tf.reduce_mean(tf.divide(1, 1 + tf.square(diffs)))

    diff = centers_batch - features
    unique_label, unique_idx, unique_count = tf.unique_with_counts(labels)
    appear_times = tf.gather(unique_count, unique_idx)
    appear_times = tf.reshape(appear_times, [-1, 1])

    diff = tf.divide(diff, tf.cast((1 + appear_times), tf.float32))
    diff = CENTER_LOSS_ALPHA * diff

    centers_update_op = tf.scatter_sub(centers, labels, diff)  # diff is used to get updated centers.

    # combo_loss = value_factor * loss + new_factor * loss2
    combo_loss = FLAGS.lambda_c * loss + FLAGS.lambda_g * loss2

    return combo_loss, centers_update_op 

def _has_enough_pixels_of_each_object_in_first_frame(
    label, decoder_output_stride):
  """Checks if for each object (incl. background) enough pixels are visible.

  During test time, we will usually not see a reference frame in which only
  very few pixels of one object are visible. These cases can be problematic
  during training, especially if more than the 1-nearest neighbor is used.
  That's why this function can be used to detect and filter these cases.

  Args:
    label: Label tensor of shape [num_frames, height, width, 1].
    decoder_output_stride: Integer, the stride of the decoder output.

  Returns:
    Boolean, whether the labels have enough pixels of each object in the first
      frame.
  """
  h, w = train_utils.resolve_shape(label)[1:3]
  h_sub = model.scale_dimension(h, 1.0 / decoder_output_stride)
  w_sub = model.scale_dimension(w, 1.0 / decoder_output_stride)
  label_downscaled = tf.squeeze(
      tf.image.resize_nearest_neighbor(label[0, tf.newaxis], [h_sub, w_sub],
                                       align_corners=True), axis=0)
  _, _, counts = tf.unique_with_counts(
      tf.reshape(label_downscaled, [-1]))
  has_enough_pixels_per_object = tf.reduce_all(
      tf.greater_equal(counts, MIN_LABEL_COUNT))
  return has_enough_pixels_per_object 

def center_loss(features, labels, num_classes, alpha=0.5, updates_collections=tf.GraphKeys.UPDATE_OPS, scope=None):
    # modified from https://github.com/EncodeTS/TensorFlow_Center_Loss/blob/master/center_loss.py

    assert features.shape.ndims == 2, 'The rank of `features` should be 2!'
    assert 0 <= alpha <= 1, '`alpha` should be in [0, 1]!'

    with tf.variable_scope(scope, 'center_loss', [features, labels]):
        centers = tf.get_variable('centers', shape=[num_classes, features.get_shape()[-1]], dtype=tf.float32,
                                  initializer=tf.constant_initializer(0), trainable=False)

        centers_batch = tf.gather(centers, labels)
        diff = centers_batch - features
        _, unique_idx, unique_count = tf.unique_with_counts(labels)
        appear_times = tf.gather(unique_count, unique_idx)
        appear_times = tf.reshape(appear_times, [-1, 1])
        diff = diff / tf.cast((1 + appear_times), tf.float32)
        diff = alpha * diff
        update_centers = tf.scatter_sub(centers, labels, diff)

        center_loss = 0.5 * tf.reduce_mean(tf.reduce_sum((centers_batch - features)**2, axis=-1))

        if updates_collections is None:
            with tf.control_dependencies([update_centers]):
                center_loss = tf.identity(center_loss)
        else:
            tf.add_to_collections(updates_collections, update_centers)

    return center_loss, centers 

def call(self, inputs, mask=None):
        features, feature_graph_index = inputs
        feature_graph_index = tf.reshape(feature_graph_index, (-1,))
        _, _, count = tf.unique_with_counts(feature_graph_index)
        m = kb.dot(features, self.m_weight)
        if self.use_bias:
            m += self.m_bias

        self.h = tf.zeros(tf.stack(
            [tf.shape(input=features)[0], tf.shape(input=count)[0], self.n_hidden]))
        self.c = tf.zeros(tf.stack(
            [tf.shape(input=features)[0], tf.shape(input=count)[0], self.n_hidden]))
        q_star = tf.zeros(tf.stack(
            [tf.shape(input=features)[0], tf.shape(input=count)[0], 2 * self.n_hidden]))
        for i in range(self.T):
            self.h, c = self._lstm(q_star, self.c)
            e_i_t = tf.reduce_sum(
                input_tensor=m * repeat_with_index(self.h, feature_graph_index), axis=-1)
            exp = tf.exp(e_i_t)
            # print('exp shape ', exp.shape)
            seg_sum = tf.transpose(
                a=tf.math.segment_sum(
                    tf.transpose(a=exp, perm=[1, 0]),
                    feature_graph_index),
                perm=[1, 0])
            seg_sum = tf.expand_dims(seg_sum, axis=-1)
            # print('seg_sum shape', seg_sum.shape)
            interm = repeat_with_index(seg_sum, feature_graph_index)
            # print('interm shape', interm.shape)
            a_i_t = exp / interm[..., 0]
            # print(a_i_t.shape)
            r_t = tf.transpose(a=tf.math.segment_sum(
                tf.transpose(a=tf.multiply(m, a_i_t[:, :, None]), perm=[1, 0, 2]),
                feature_graph_index), perm=[1, 0, 2])
            q_star = kb.concatenate([self.h, r_t], axis=-1)
        return q_star 

def build(self, predictions, targets, inputs=None):
        """ Prints the number of each kind of prediction """
        self.built = True
        pshape = predictions.get_shape()
        self.inner_metric.build(predictions, targets, inputs)

        with tf.name_scope(self.name):
            if len(pshape) == 1 or (len(pshape) == 2 and int(pshape[1]) == 1):
                self.name = self.name or "binary_prediction_counts"
                y, idx, count = tf.unique_with_counts(tf.argmax(predictions))
                self.tensor = tf.Print(self.inner_metric, [y, count], name=self.inner_metric.name)
            else:
                self.name = self.name or "categorical_prediction_counts"
                y, idx, count = tf.unique_with_counts(tf.argmax(predictions, dimension=1))
                self.tensor = tf.Print(self.inner_metric.tensor, [y, count], name=self.inner_metric.name) 

def get_center_loss(features, labels, alpha, num_classes):
    """
    Arguments:
        features: Tensor,shape [batch_size, feature_length].
        labels: Tensor,shape [batch_size].#not the one hot label
        alpha:  center upgrade learning rate
        num_classes: how many classes. 
    
    Return：
        loss: Tensor,
        centers: Tensor
        centers_update_op:
    """
    len_features = features.get_shape()[1]
    centers = tf.get_variable('centers', [num_classes, len_features], dtype=tf.float32,
        initializer=tf.constant_initializer(0), trainable=False)
    labels = tf.reshape(labels, [-1])
    centers_batch = tf.gather(centers, labels)
    loss = tf.nn.l2_loss(features - centers_batch)
    diff = centers_batch - features
    unique_label, unique_idx, unique_count = tf.unique_with_counts(labels)
    appear_times = tf.gather(unique_count, unique_idx)
    appear_times = tf.reshape(appear_times, [-1, 1])
    diff = diff / tf.cast((1 + appear_times), tf.float32)
    diff = alpha * diff
    centers_update_op = tf.scatter_sub(centers, labels, diff)
    # need to update after every epoch, the key is to update the center of the classes. 

    return loss, centers, centers_update_op 

def reduce_batch_weighted_counts(x, weights=None):
  """Performs batch-wise reduction to produce (possibly weighted) counts.

  Args:
    x: Input `Tensor`.
    weights: (Optional) Weights input `Tensor`.

  Returns:
    a named tuple of...
      The unique values in x
      The sum of the weights for each unique value in x if weights are provided,
        else None
  """
  if isinstance(x, tf.SparseTensor):
    x = x.values
  if weights is None:
    # TODO(b/112916494): Always do batch wise reduction once possible.

    return ReducedBatchWeightedCounts(tf.reshape(x, [-1]), None, None, None)
  # TODO(b/134075780): Revisit expected weights shape when input is sparse.
  x, weights = assert_same_shape(x, weights)
  weights = tf.reshape(weights, [-1])
  x = tf.reshape(x, [-1])
  unique_x_values, unique_idx, _ = tf.unique_with_counts(x, out_idx=tf.int64)
  summed_weights_per_x = tf.math.unsorted_segment_sum(
      weights, unique_idx, tf.size(input=unique_x_values))
  return ReducedBatchWeightedCounts(unique_x_values, summed_weights_per_x, None,
                                    None) 

def testString(self):
    indx = np.random.randint(65, high=122, size=7000)
    x = [chr(i) for i in indx]

    with self.test_session() as sess:
      y, idx, count = tf.unique_with_counts(x)
      tf_y, tf_idx, tf_count = sess.run([y, idx, count])

    self.assertEqual(len(x), len(tf_idx))
    self.assertEqual(len(tf_y), len(np.unique(x)))
    for i in range(len(x)):
      self.assertEqual(x[i], tf_y[tf_idx[i]].decode('ascii'))
    for value, count in zip(tf_y, tf_count):
      v = [1 if x[i] == value.decode('ascii') else 0 for i in range(7000)]
      self.assertEqual(count, sum(v)) 

def testInt32(self):
    x = np.random.randint(2, high=10, size=7000)
    with self.test_session() as sess:
      y, idx, count = tf.unique_with_counts(x)
      tf_y, tf_idx, tf_count = sess.run([y, idx, count])

    self.assertEqual(len(x), len(tf_idx))
    self.assertEqual(len(tf_y), len(np.unique(x)))
    for i in range(len(x)):
      self.assertEqual(x[i], tf_y[tf_idx[i]])
    for value, count in zip(tf_y, tf_count):
      self.assertEqual(count, np.sum(x == value)) 

def get_top_elements(list_of_elements, max_user_contribution):
  """Gets the top max_user_contribution words from the input list.

  Note that the returned set of top words will not necessarily be sorted.

  Args:
    list_of_elements: A tensor containing a list of elements.
    max_user_contribution: The maximum number of elements to keep.

  Returns:
    A tensor of a list of strings.
    If the total number of unique words is less than or equal to
    max_user_contribution, returns the set of unique words.
  """
  words, _, counts = tf.unique_with_counts(list_of_elements)
  if tf.size(words) > max_user_contribution:
    # This logic is influenced by the focus on global heavy hitters and
    # thus implements clipping by chopping the tail of the distribution
    # of the words as present on a single client. Another option could
    # be to provide pick max_words_per_user random words out of the unique
    # words present locally.
    top_indices = tf.argsort(
        counts, axis=-1, direction='DESCENDING')[:max_user_contribution]
    top_words = tf.gather(words, top_indices)
    return top_words
  return words 

def center_loss_v2(config, features, labels, centers=None, **kargs):
	alpha = config.alpha
	num_classes = config.num_classes
	with tf.variable_scope(config.scope+"_center_loss"):
		print("==center loss==")
		len_features = features.get_shape()[1]
		if not centers:
			centers = tf.get_variable('centers', 
							[num_classes, len_features], 
							dtype=tf.float32,
							initializer=tf.contrib.layers.xavier_initializer(),
							trainable=False)
			print("==add center parameters==")
	 
		centers_batch = tf.gather(centers, labels)

		loss = tf.nn.l2_loss(features - centers_batch)
	 
		diff = centers_batch - features
	 
		unique_label, unique_idx, unique_count = tf.unique_with_counts(labels)
		appear_times = tf.gather(unique_count, unique_idx)
		appear_times = tf.reshape(appear_times, [-1, 1])
	 
		diff = diff / tf.cast((1 + appear_times), tf.float32)
		diff = alpha * diff

		centers_update_op = tf.scatter_sub(centers, labels, diff)

		tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, centers_update_op)
		
		return loss, centers 

def weighted_loss_ratio(config, losses, labels, ratio_weight):
	unique_label, unique_idx, unique_count = tf.unique_with_counts(labels)
	appear_times = tf.gather(unique_count, unique_idx)
	# appear_times = tf.reshape(appear_times, [-1, 1])

	weighted_loss = losses * ratio_weight
	weighted_loss = weighted_loss / tf.cast((EPSILON+appear_times), tf.float32)

	return weighted_loss, None 

def reduce_batch_count_mean_and_var_per_key(x, key, reduce_instance_dims):
  """Computes per-key element count, mean and var for the given tensor.

  Args:
    x: A `Tensor` or `SparseTensor`.
    key: A `Tensor` or `SparseTensor` (cannot be None).
        Must meet one of the following conditions:
        1. Both x and key are dense,
        2. Both x and key are sparse and `key` must exactly match `x` in
        everything except values,
        3. The axis=1 index of each x matches its index of dense key.
    reduce_instance_dims: A bool, if True - collapses the batch and instance
        dimensions to arrive at a single scalar output. Otherwise, only
        collapses the batch dimension and outputs a `Tensor` of the same shape
        as the input. Not supported for `SparseTensor`s.

  Returns:
    A 4-tuple containing the `Tensor`s (key_vocab, count, mean, var).
  """

  if isinstance(x, tf.SparseTensor):
    if not reduce_instance_dims:
      raise NotImplementedError(
          'Mean and var per key only support reduced dims for SparseTensors')

  x, key = _validate_and_get_dense_value_key_inputs(x, key)

  unique = tf.unique_with_counts(key, out_idx=tf.int64)
  x_count = unique.count
  x_count = tf.cast(x_count, x.dtype)
  if not reduce_instance_dims:
    x_count = tf.tile(tf.expand_dims(x_count, axis=-1), [1, x.shape[1]])

  if reduce_instance_dims:
    sums = tf.reduce_sum(x, axis=1) if x.get_shape().ndims != 1 else x
    sums = tf.math.unsorted_segment_sum(sums, unique.idx, tf.size(unique.y))
  else:
    sums = tf.math.unsorted_segment_sum(x, unique.idx, tf.size(unique.y))

  means = tf.cast(sums, x.dtype) / x_count
  sum_sqs = tf.math.unsorted_segment_sum(tf.square(x),
                                         unique.idx,
                                         tf.size(input=unique.y))
  if sum_sqs.get_shape().ndims != 1 and reduce_instance_dims:
    sum_sqs = tf.reduce_sum(sum_sqs, axis=1)

  variances = sum_sqs / x_count - tf.square(means)

  return unique.y, x_count, means, variances


# Code for serializing and example proto 

def reduce_batch_minus_min_and_max_per_key(x, key):
  """Computes the -min and max of a tensor x.

  Args:
    x: A `tf.Tensor` or `SparseTensor`.
    key: A `Tensor` or `SparseTensor`.
        Must meet one of the following conditions:
        1. Both x and key are dense,
        2. Both x and key are sparse and `key` must exactly match `x` in
        everything except values,
        3. The axis=1 index of each x matches its index of dense key.
  Returns:
    A 3-tuple containing the `Tensor`s (key_vocab, min_per_key, max_per_key).
  """
  output_dtype = x.dtype

  if x.dtype == tf.uint8 or x.dtype == tf.uint16:
    x = tf.cast(x, tf.int32)

  elif x.dtype == tf.uint32 or x.dtype == tf.uint64:
    raise TypeError('Tensor type %r is not supported' % x.dtype)

  x, key = _validate_and_get_dense_value_key_inputs(x, key)

  def get_batch_max_per_key(tensor, key_uniques, dtype):  # pylint: disable=missing-docstring
    if tensor.get_shape().ndims < 2:
      row_maxes = tensor
    else:
      row_maxes = tf.reduce_max(
          tensor, axis=tf.range(1, tensor.get_shape().ndims))
    batch_max = tf.math.unsorted_segment_max(
        row_maxes, key_uniques.idx, tf.size(input=key_uniques.y))

    # TODO(b/112309021): Remove workaround once tf.reduce_max of a tensor of all
    # NaNs produces -inf.
    return _inf_to_nan(batch_max, dtype)

  unique = tf.unique_with_counts(key, out_idx=tf.int64)
  x_batch_maxes = get_batch_max_per_key(x, unique, output_dtype)
  x_batch_minus_mins = get_batch_max_per_key(-x, unique, output_dtype)

  x_batch_minus_mins, x_batch_maxes = assert_same_shape(x_batch_minus_mins,
                                                        x_batch_maxes)

  return (unique.y, x_batch_minus_mins, x_batch_maxes) 

def compute_spans(start_scores, end_scores, answer2support, is_eval, support2question,
                  topk=1, max_span_size=10000, correct_start=None):
    max_support_length = tf.shape(start_scores)[1]
    _, _, num_doc_per_question = tf.unique_with_counts(support2question)
    offsets = tf.cumsum(num_doc_per_question, exclusive=True)
    doc_idx_for_support = tf.range(tf.shape(support2question)[0]) - tf.gather(offsets, support2question)

    def train():
        gathered_end_scores = tf.gather(end_scores, answer2support)
        gathered_start_scores = tf.gather(start_scores, answer2support)

        if correct_start is not None:
            # assuming we know the correct start we only consider ends after that
            left_mask = misc.mask_for_lengths(tf.cast(correct_start, tf.int32), max_support_length, mask_right=False)
            gathered_end_scores = gathered_end_scores + left_mask

        predicted_start_pointer = tf.argmax(gathered_start_scores, axis=1, output_type=tf.int32)
        predicted_end_pointer = tf.argmax(gathered_end_scores, axis=1, output_type=tf.int32)

        return (start_scores, end_scores,
                tf.gather(doc_idx_for_support, answer2support), predicted_start_pointer, predicted_end_pointer)

    def eval():
        # we collect spans for top k starts and top k ends and select the top k from those top 2k
        doc_idx1, start_pointer1, end_pointer1, span_score1 = _get_top_k(
            start_scores, end_scores, topk, max_span_size, support2question)
        doc_idx2, end_pointer2, start_pointer2, span_score2 = _get_top_k(
            end_scores, start_scores, topk, -max_span_size, support2question)

        doc_idx = tf.concat([doc_idx1, doc_idx2], 1)
        start_pointer = tf.concat([start_pointer1, start_pointer2], 1)
        end_pointer = tf.concat([end_pointer1, end_pointer2], 1)
        span_score = tf.concat([span_score1, span_score2], 1)

        _, idx = tf.nn.top_k(span_score, topk)

        r = tf.range(tf.shape(span_score)[0], dtype=tf.int32)
        r = tf.reshape(tf.tile(tf.expand_dims(r, 1), [1, topk]), [-1, 1])

        idx = tf.concat([r, tf.reshape(idx, [-1, 1])], 1)
        doc_idx = tf.gather_nd(doc_idx, idx)
        start_pointer = tf.gather_nd(start_pointer, idx)
        end_pointer = tf.gather_nd(end_pointer, idx)

        return start_scores, end_scores, tf.gather(doc_idx_for_support, doc_idx), start_pointer, end_pointer

    return tf.cond(is_eval, eval, train) 

def discriminative_loss_single_multicate(sem_label, prediction, correct_label, feature_dim,
                               delta_v, delta_d, param_var, param_dist, param_reg):
    ''' Discriminative loss for a single prediction/label pair.
    :param sem_label: semantic label
    :param prediction: inference of network
    :param correct_label: instance label
    :feature_dim: feature dimension of prediction
    :param label_shape: shape of label
    :param delta_v: cutoff variance distance
    :param delta_d: curoff cluster distance
    :param param_var: weight for intra cluster variance
    :param param_dist: weight for inter cluster distances
    :param param_reg: weight regularization
    '''
    unique_sem_label, unique_id, counts = tf.unique_with_counts(sem_label)
    num_sems = tf.size(unique_sem_label)

    def cond(i, ns, unique_id, pred, ins_label, out_loss, out_var, out_dist, out_reg):
        return tf.less(i, num_sems)

    def body(i, ns, unique_id, pred, ins_label, out_loss, out_var, out_dist, out_reg):
        inds = tf.equal(i, unique_id)
        cur_pred = tf.boolean_mask(prediction, inds)
        cur_label = tf.boolean_mask(correct_label, inds)
        cur_discr_loss, cur_l_var, cur_l_dist, cur_l_reg = discriminative_loss_single(cur_pred, cur_label, feature_dim,
                            delta_v, delta_d, param_var, param_dist, param_reg)
        out_loss = out_loss.write(i, cur_discr_loss)
        out_var = out_var.write(i, cur_l_var)
        out_dist = out_dist.write(i, cur_l_dist)
        out_reg = out_reg.write(i, cur_l_reg)

        return i + 1, ns, unique_id, pred, ins_label, out_loss, out_var, out_dist, out_reg

    output_ta_loss = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)
    output_ta_var = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)
    output_ta_dist = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)
    output_ta_reg = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)

    loop = [0, num_sems, unique_id, prediction, correct_label, output_ta_loss, output_ta_var, output_ta_dist, output_ta_reg]
    _, _, _, _, _, out_loss_op, out_var_op, out_dist_op, out_reg_op = tf.while_loop(cond, body, loop)

    out_loss_op = out_loss_op.stack()
    out_var_op = out_var_op.stack()
    out_dist_op = out_dist_op.stack()
    out_reg_op = out_reg_op.stack()

    disc_loss = tf.reduce_sum(out_loss_op)
    l_var = tf.reduce_sum(out_var_op)
    l_dist = tf.reduce_sum(out_dist_op)
    l_reg = tf.reduce_sum(out_reg_op)

    return disc_loss, l_var, l_dist, l_reg 

def get_center_loss_tf(self, features):
        """获取center loss及center的更新op

            Arguments:
                features: Tensor,表征样本特征,一般使用某个fc层的输出,shape应该为[batch_size, feature_length].
                labels: Tensor,表征样本label,非one-hot编码,shape应为[batch_size].
                alpha: 0-1之间的数字,控制样本类别中心的学习率,细节参考原文.
                num_classes: 整数,表明总共有多少个类别,网络分类输出有多少个神经元这里就取多少.

            Return：
                loss: Tensor,可与softmax loss相加作为总的loss进行优化.
                centers: Tensor,存储样本中心值的Tensor，仅查看样本中心存储的具体数值时有用.
                centers_update_op: op,用于更新样本中心的op，在训练时需要同时运行该op，否则样本中心不会更新
        """
        # 获取特征的维数，例如256维
        len_features = features.get_shape()[1]
        # 建立一个Variable,shape为[num_classes, len_features]，用于存储整个网络的样本中心，
        # 设置trainable=False是因为样本中心不是由梯度进行更新的
        centers = tf.get_variable('centers',
                                  shape=[self.num_classes, len_features],
                                  dtype=tf.float32,
                                  initializer=tf.constant_initializer(0),
                                  trainable=False)
        # 将label展开为一维的，输入如果已经是一维的，则该动作其实无必要
        labels = tf.reshape(self.labels, [-1])

        # 根据样本label,获取mini-batch中每一个样本对应的中心值
        centers_batch = tf.gather(centers, labels)
        # 计算loss
        loss = tf.nn.l2_loss(features - centers_batch)

        # 当前mini-batch的特征值与它们对应的中心值之间的差
        diff = centers_batch - features

        # 获取mini-batch中同一类别样本出现的次数,了解原理请参考原文公式(4)
        unique_label, unique_idx, unique_count = tf.unique_with_counts(labels)
        appear_times = tf.gather(unique_count, unique_idx)
        appear_times = tf.reshape(appear_times, [-1, 1])

        diff = diff / tf.cast((1 + appear_times), tf.float32)
        diff = self.alpha * diff

        centers_update_op = tf.scatter_sub(centers, labels, diff)

        return loss, centers, centers_update_op