Stream compaction recipe

General Discussion
#1

Hi All,

My first post on the TF Forum, so hello!

What would be the recommended recipe for compacting a stream in TF? For example, suppose I have

foo = tf.constant([1, 2, 3, 4, 5])
mask = tf.constant([0, 1, 0, 1,1])

and I wish to gather the elements of foo corresponding to elements of mask equal to 1 such that

compact(foo, mask) => [2, 4, 5]

In C++, thrust::copy_if allows this easily. What would be the TF way? tf.dynamic_partition might do this, but seems not quite the right function…

Many thanks,

Chris

#3

Are you looking for this?

https://www.tensorflow.org/api_docs/python/tf/boolean_mask

#5

Sure, thanks! But is there an XLA-able option? For example some kind of sort, that piles up the required values into the first k indices of a tensor, and returns k?

Chris

#7

But is there an XLA-able option? For example some kind of sort, that piles up the required values into the first k indices of a tensor, and returns k?

Is this another question?

#8

Hmm, I stand corrected. I remembered that tf.boolean_mask could not be jit-compiled by XLA. However, the following works:

x = tf.constant([1, 0, 2, 3, 0, 0, 4])
def compact(x):
   msk = x > 0
   return tf.boolean_mask(x, msk)
xla_compact = tf.function(xla_compact, jit_compile=True)

compact(x)
# <tf.Tensor: shape=(4,), dtype=int32, numpy=array([1, 2, 3, 4], dtype=int32)>
xla_compact(x)
# <tf.Tensor: shape=(4,), dtype=int32, numpy=array([1, 2, 3, 4], dtype=int32)>

All good. Many thanks!

1 Like