Suppose you have such a tensor
a = [0, 0, 0, 1, 1, 2]
and want to know for each i
in a
it is the n
th i
in a
, that is, to get
b = [0, 1, 2, 0, 1, 0]
Here's the solution:
count_sum = tf.cumsum(tf.ones_like(a))
count_min = tf.segment_min(count_sum, a)
b = count_sum - tf.gather(count_min, a)
Now let a
be unsorted:
a = [0, 1, 0, 2, 1, 2, 1, 2, 2, 2]
Then we need to sort a
first
a = tf.constant([0,1,0,2,1,2,1,2,2,2])
a_args = tf.contrib.framework.argsort(a, stable=True)
a_sort = tf.gather(a, a_args)
count_sum = tf.cumsum(tf.ones_like(a_sort))
count_min = tf.segment_min(count_sum, a_sort)
b_sort = count_sum - tf.gather(count_min, a_sort)
b = tf.unsorted_segment_sum(b_sort, a_args, b_sort.shape[0])