I have a tensor of allowed items
index = tf.constant([61, 215, 23, 18, 241, 125])
and need to remove items from input sequence batches that are not in index.
seq = tf.constant(
[
[ 18, 241, 0, 0],
[125, 61, 23, 241],
[ 23, 92, 18, 0],
[ 5, 61, 215, 18],
]
)
After the calculation in this case I need
result_needed = tf.constant(
[
[ 18, 241, 0, 0],
[125, 61, 23, 241],
[ 23, 18, 0, 0],
[ 61, 215, 18, 0],
]
)
I cannot do this in Python because this calculation happens during predictions. Also note that while item IDs here are small, solution needs to deal with numbers from 1 to 2^40.
Answer
After some serious pondering time, I came up with the following:
idx_range = tf.reshape(tf.range(seq.shape[-2]), [-1, 1])
idx_tile = tf.tile(idx_range, [1, seq.shape[-2].value])
idx_flat = tf.reshape(idx_tile, [-1])
truth_value = tf.equal(index, tf.expand_dims(seq, -1))
one_hot = tf.to_float(truth_value)
ones = tf.nn.top_k(tf.reduce_sum(one_hot, -1), seq.shape[-1]).indices
ones_flat = tf.reshape(ones, [-1])
ones_idx = tf.reshape(
tf.stack([idx_flat, ones_flat], axis=1),
tf.concat([seq.shape, [2]], axis=0)
)
tf.gather_nd(seq, ones_idx)
This is not exactly what I said I needed, but actually got me close enough. Instead of the output replacing the blacklisted items with 0, it moves them to the end. If you needed them gone, I'm sure there's a method to remove them, but I'm not looking into it. Apologies.