Find Unique values in a 2D Tensor using Tensorflow

Question

tf.unique currently only works on 1D tensors. How can I find unique values in a 2D tensor.

ip=tf.constant([[1,2,1],[3,4,1],[5,6,1],[1,2,1]])
#op should be  = [[1,2,1],[3,4,1],[5,6,1]]

Answer 1

Following function takes 2D tensor as IP and will return Unique 2D Tensor

def tf_unique_2d(self, x):
    x_shape = tf.shape(x)  # (3,2)
    x1 = tf.tile(x, [1, x_shape[0]])  # [[1,2],[1,2],[1,2],[3,4],[3,4],[3,4]..]
    x2 = tf.tile(x, [x_shape[0], 1])  # [[1,2],[1,2],[1,2],[3,4],[3,4],[3,4]..]

    x1_2 = tf.reshape(x1, [x_shape[0] * x_shape[0], x_shape[1]])
    x2_2 = tf.reshape(x2, [x_shape[0] * x_shape[0], x_shape[1]])
    cond = tf.reduce_all(tf.equal(x1_2, x2_2), axis=1)
    cond = tf.reshape(cond, [x_shape[0], x_shape[0]])  # reshaping cond to match x1_2 & x2_2
    cond_shape = tf.shape(cond)
    cond_cast = tf.cast(cond, tf.int32)  # convertin condition boolean to int
    cond_zeros = tf.zeros(cond_shape, tf.int32)  # replicating condition tensor into all 0's

    # CREATING RANGE TENSOR
    r = tf.range(x_shape[0])
    r = tf.add(tf.tile(r, [x_shape[0]]), 1)
    r = tf.reshape(r, [x_shape[0], x_shape[0]])

    # converting TRUE=1 FALSE=MAX(index)+1 (which is invalid by default) so when we take min it wont get selected & in end we will only take values <max(indx).
    f1 = tf.multiply(tf.ones(cond_shape, tf.int32), x_shape[0] + 1)
    f2 = tf.ones(cond_shape, tf.int32)
    cond_cast2 = tf.where(tf.equal(cond_cast, cond_zeros), f1, f2)  # if false make it max_index+1 else keep it 1

    # multiply range with new int boolean mask
    r_cond_mul = tf.multiply(r, cond_cast2)
    r_cond_mul2 = tf.reduce_min(r_cond_mul, axis=1)
    r_cond_mul3, unique_idx = tf.unique(r_cond_mul2)
    r_cond_mul4 = tf.subtract(r_cond_mul3, 1)

    # get actual values from unique indexes
    op = tf.gather(x, r_cond_mul4)

    return (op)

Answer 2

The easiest way to do this is to first reshape your tensor to be a 1D tensor, then use tf.unique (which currently only accepts 1D tensors, not sure why!) or tf.unique_with_counts (depending on whether you also need the counts of the unique elements), and then, if you also need the indices of the unique values to have the same shape as the original tensor, you can reshape also the indices. The following TensorFlow 2.1 program illustrates this.

import tensorflow as tf # TF 2.1


def get_uniques(t):
    t1d = tf.reshape(t, shape=(-1,))
    # or tf.unique, if you don't need counts
    uniques, idx, counts = tf.unique_with_counts(t1d) 
    return uniques, tf.reshape(idx, shape=tf.shape(t)), counts


if __name__ == '__main__':
    my_tensor = tf.constant([[1, 2, 1], [3, 4, 1], [5, 6, 1], [1, 2, 1]])

    uniques, idx, counts = get_uniques(my_tensor)

    tf.print("tf.shape(uniques) =", tf.shape(uniques))
    tf.print("tf.shape(idx) =", tf.shape(idx))
    tf.print("tf.shape(counts) =", tf.shape(counts))
    tf.print("uniques =", uniques)

Answer 3

Answer in April 2022

Since version 2.2, Tensorflow has supported find unique 1-D tensor in a 2D tensor using tf.raw_ops.UniqueV2.

>>> import tensorflow as tf
>>> a = tf.constant([[1, 2, 1], [3, 4, 1], [5, 6, 1], [1, 2, 1]])
>>> y, idx = tf.raw_ops.UniqueV2(x=a, axis=[0])
>>> y
<tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 1],
       [3, 4, 1],
       [5, 6, 1]], dtype=int32)>
>>> idx
<tf.Tensor: shape=(4,), dtype=int32, numpy=array([0, 1, 2, 0], dtype=int32)>

Check the docs for more details.