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

Bhack
February 3, 2022, 7:51am
3
Are you looking for this?

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

Bhack
February 3, 2022, 1:08pm
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?

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