I have a ragged tensor:
tf.ragged.constant([[[17712], [16753], [11850], [13028], [10155], [15734, 15938], [126], [10135], [17665]]], dtype=tf.int32)
I would like to set the value of elements in rows with a length greater than 1 to a particular value. For example:
tf.ragged.constant([[[17712], [16753], [11850], [13028], [10155], [15734, 0], [126], [10135], [17665]]], dtype=tf.int32)
How can I express such a transformation in Tensorflow?
Ragged tensors always make things trickier than they should be, but here is one possible implementation of that:
import tensorflow as tf
# Using an intermediate NumPy array avoids having the second dimension as ragged
a = tf.ragged.constant([[[17712], [16753], [11850], [13028], [10155], [15734, 15938],
[126], [10135], [17665]]], dtype=tf.int32)
# Index from which values are replaced
replace_from_idx = 1
# Replacement value
new_value = 0
# Get size of each element in the last dimension
s = a.row_lengths(axis=-1)
# Make ragged ranges
r = tf.ragged.range(s.flat_values)
# Un-flatten
r = tf.RaggedTensor.from_row_lengths(r, a.row_lengths(1))
# Replace values
m = tf.dtypes.cast(r < replace_from_idx, a.dtype)
out = a * m + new_value * (1 - m)
print(out.to_list())
# [[[17712], [16753], [11850], [13028], [10155], [15734, 0], [126], [10135], [17665]]]