Tensorflow - Conditionally assign value to a tensor

Question

I have 2 tensors.

t1 = tf.constant([b'hi', b'#hh', b'hello', 'there', '#ii'], dtype=tf.string)

t2 = tf.constant([1,2,3], dtype = tf.int64)

How can I replace the elements, initialized with #, in t1 to be 0 and all other elements in t1 to be the numbers in t2?

Specifically, the output should be [1, 0, 2, 3, 0].

Thanks,

Answer 1

You can do the following. This should work with both TF1 and TF2 (Tested). In TF1 just do a sess.run(res) to get the result.

t1 = tf.constant(['hi', '#hh', 'hello', 'there', '#ii'], dtype=tf.string)
t2 = tf.constant([1,2,3], dtype = tf.int32)

# Create mask of the ids we want to change
mask = tf.logical_not(tf.strings.regex_full_match(t1,'^#.*'))
idx = tf.cast(mask, dtype=tf.int32)
# Get the cumulative sum [1, 1, 2, 3, 3]
cumsum = tf.cumsum(idx)
# Do a where and replace the unwanted ids with zeros
res = tf.where(mask, y=tf.zeros_like(t1, dtype=tf.int32), x=cumsum)