TensorFlow - return distinct sub-tensors of multidimensional tensor

Question

In TensorFlow, the tf.unique function can be used to return the distinct elements of a 1-dimensional Tensor. How can I get the distinct sub-Tensors along the axis 0 of a higher-dimensional Tensor? For example, given the following Tensor, the desired distinct function would return the specified result:

input = tf.constant([
    [0,3],
    [0,1],
    [0,4],
    [0,1],
    [1,5],
    [3,9],
    [3,2],
    [3,6],
    [3,5],
    [3,3]])

distinct(input) == tf.constant([
    [0,3],
    [0,1],
    [0,4],
    [1,5],
    [3,9],
    [3,2],
    [3,6],
    [3,5],
    [3,3]])

How can the distinct multidimensional elements be generated for Tensors of any number of dimensions?

Answer 1

Without preserving Order

You can use tf.py_function and call np.unique to return unique multidimensional tensors along axis=0. Note that this finds the unique rows but does not preserve the order.

def distinct(a):
    _a =  np.unique(a, axis=0)
    return _a

>> input = tf.constant([
[0,3],
[0,1],
[0,4],
[0,1],
[1,5],
[3,9],
[3,2],
[3,6],
[3,5],
[3,3]])

>> tf.py_function(distinct, [input], tf.int32)
<tf.Tensor: id=940, shape=(9, 2), dtype=int32, numpy=
array([[0, 1],
   [0, 3],
   [0, 4],
   [1, 5],
   [3, 2],
   [3, 3],
   [3, 5],
   [3, 6],
   [3, 9]], dtype=int32)>

With Orders preserved

def distinct_with_order_preserved(a):
    _a = a.numpy()
    return pd.DataFrame(_a).drop_duplicates().values

>> tf.py_function(distinct_with_order_preserved, [input], tf.int32)
<tf.Tensor: id=950, shape=(9, 2), dtype=int32, numpy=
array([[0, 3],
   [0, 1],
   [0, 4],
   [1, 5],
   [3, 9],
   [3, 2],
   [3, 6],
   [3, 5],
   [3, 3]], dtype=int32)>