我如何取出（或切片）秩2张量中的元素，其第一个元素是唯一的？

Question

My title might be ambiguous due to my awkward English. But I mean this: suppose i have a tensor a like this:

array([[1, 2, 3],
       [2, 2, 3],
       [2, 2, 4],
       [3, 2, 3],
       [4, 2, 3]], dtype=int32)

the 'first column' of this tensor could contain duplicate elements (eg [1, 2, 2, 3, 4] or [1, 1, 2, 3, 3, 4, 5, 5]), and which element is duplicated is not known beforehand.

and i wanna take out a tensor this:

array([[1, 2, 3],
       [2, 2, 3],
       [3, 2, 3],
       [4, 2, 3]], dtype=int32)

as u can see, I take out the rows whose first element is a unique element in the column of a .

I first wanted to use the function tf.unique() . BUT the idx value returned by it doesn't indicate the first index of each value of output tensor in the original tensor.

tf.unique() works like this:

# tensor 'x' is [1, 1, 2, 3, 3, 3, 7, 8, 8]
y, idx = tf.unique(x)
y ==> [1, 2, 3, 7, 8]
idx ==> [0, 0, 1, 2, 2, 2, 3, 4, 4]

The function tf.unique(x, name=None) finds the unique elements in a 1-D tensor. And it now returns two value: y and idx . y contains all of the unique elements of x sorted inthe same order that they occur in x . idx contains the index of each value of x in the unique output y .

How I wish it has a third return value which contains the first index of each value of y in the original tensor x is also needed. It might work like this:

# tensor 'x' is [1, 1, 2, 3, 3, 3, 7, 8, 8]
y, idx, idx_ori  = tf.unique(x)
y ==> [1, 2, 3, 7, 8]
idx ==> [0, 0, 1, 2, 2, 2, 3, 4, 4]
idx_ori ==> [0, 2, 3, 6, 7]

Just like its equivalent in Numpy does:

array 'x' is [1, 1, 2, 3, 3, 3, 7, 8, 8]
y, idx_ori = np.unique(x, return_index=True)
y ==> [1, 2, 3, 7, 8]
idx_ori ==> [0, 2, 3, 6, 7]

IF i have this idx_ori , i can solve my problem by tf.gather() :

_, _1, idx_ori = tf.unique(a[:, 0])
result = tf.gather(a, idx_ori)

Any idea to workaround this problem? or any idea to get this indices that i want.

PS I know my description is tediously long ... :-p

Answer 1

This is a bit gross, but you could do:

print a
y, idx = tf.unique(a[:,0])
z = tf.one_hot(idx, tf.shape(y)[0])
s = tf.cumsum(z)
e = tf.equal(s, 1)  # only seen once so far
ss = tf.to_int32(e) * tf.to_int32(z) # and we equal the thing
m = tf.reduce_max(ss, reduction_indices=1)
out = tf.boolean_mask(a, tf.equal(m, 1))
sess = tf.Session()
print sess.run(out)

[[1 2 3]
 [2 2 3]
 [2 2 4]
 [3 2 3]
 [4 2 3]]
[[1 2 3]
 [2 2 3]
 [3 2 3]
 [4 2 3]]