Tensorflow remove a list of elements from a tensor

Viewed 294

I want to remove a list of possible values from a shape 1 tensor, with tensorflow code

For example

values = [1,3,5]
tensor = [1,2,205,23,5]
remove(tensor,values) -> [2,205,23]

And I would need this to be as scalable as possible related to values as they can be up to 10k for my use case. Is there any better than iterating over the 10k solution?

1 Answers

Simplest solution would be tf.sets.difference, which returns a sparse tensor with indices and values:

import tensorflow as tf

values = tf.constant([1,3, 5])
tensor = tf.constant([1,2,205,23,5])

tensor = tf.sets.difference(tf.expand_dims(tensor, axis=0), tf.expand_dims(values, axis=0)).values
tf.Tensor([  2  23 205], shape=(3,), dtype=int32)

You could also try using tf.repeat and tf.boolean_mask if you want to keep the order:

import tensorflow as tf

values = tf.constant([1, 3, 5])
tensor = tf.constant([1,2,205,23,5])

repeat_values = tf.repeat(values, repeats=tf.shape(tensor)[0], axis=0)
repeat_tensor = tf.repeat(tensor, repeats=tf.shape(values)[0], axis=0)
tensor = tf.boolean_mask(repeat_tensor, tf.not_equal(repeat_tensor, repeat_values))[::tf.shape(values)[0]]
# tf.Tensor([  2 205  23], shape=(3,), dtype=int32)

But IMO tf.sets.difference is more elegant.

Related