Suppose I have a tensor like this.

t1 = tf.constant([0,1,2,3,5,6,7,8,9,10,19,20,21,22,23,24])

I want to separate it into segments, containing the groups where the numbers are immediately adjacent. The expected output would be something like this:

t2 = [0,0,0,0,1,1,1,1,1,1,2,2,2,2,2,2]

Any idea on how to do this using tensorflow methods?

CodePudding user response：

import tensorflow as tf
t1 = tf.constant([0,1,2,3,5,6,7,8,9,10,19,20,21,22,23,24], dtype=tf.int32)

t2 = tf.roll(t1 - 1, shift=-1, axis=0)
output = tf.math.cumsum(tf.cast(t2 != t1, dtype=tf.int32), exclusive=True)

output # [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]

CodePudding user response：

One way to achieve this functionality would be like this:

diff = t1[1:] - t1[:-1]
jumps = tf.abs(diff) > 1

t2 = tf.math.cumsum(tf.cast(jumps, 'int32'))
t2 = tf.concat([[0], t2], 0)

We first take the difference between consecutive elements, then check if the absolute difference is greater than one (in which case the numbers are not adjacent). This yields an array with True at the places of the 'jumps' and False elsewhere. The abs makes sure that a jump to a lower number (i.e. [4, 5, 6, 1, 2]) is handled correctly. You can leave it out if you know that numbers will always be increasing.

Taking the cumulative sum of this array (casted to an integer type) will provide an array with the different sections represented as integers. Finally, since taking the difference reduces the length of the array by one, we prepend t2 with a 0, since the first number should belong to the first section.

One small note: this treats numbers decreasing by one (i.e. [4, 5, 6, 5, 4]) as adjacent and puts them in the same section.