Slice a dense tensor into ragged flat values

user_3015734 0
Viewed 273

I have a 2D dense tensor and a vector of lengths, which are the first K elements of each row I would like to select:

x = tf.constant([
    [1, 1, 1, 0, 0, 0],
    [2, 0, 2, 0, 0, 0],
    [1, 1, 0, 0, 0, 0],
    [3, 1, 1, 1, 1, 0],
])

seq_len = tf.constant([3, 3, 2, 5])

Desired output would be:

tf.constant([1, 1, 1, 2, 0, 2, 1, 1, 3, 1, 1, 1, 1])

I have tried going dense->ragged to get the flat values, but this does not work:

>>> tf.RaggedTensor.from_row_lengths(x, seq_len)

InvalidArgumentError: Arguments to _from_row_partition do not form a valid RaggedTensor
Condition x == y did not hold.
First 1 elements of x:
[13]
First 1 elements of y:
[4]

I cannot go dense -> sparse either since there is a zero I would like to keep in the output. Any ideas?

tensorflow python
1 Answers

You can slice a dense tensor to ragged tensor as shown below

import tensorflow as tf

x = tf.constant([
    [1, 1, 1, 0, 0, 0],
    [2, 0, 2, 0, 0, 0],
    [1, 1, 0, 0, 0, 0],
    [3, 1, 1, 1, 1, 0],
])

seq_len = tf.constant([3, 3, 2, 5])

To achieve desired output, you can use args padding

tf.RaggedTensor.from_tensor(x, padding=0)

Output:

<tf.RaggedTensor [[1, 1, 1], [2, 0, 2], [1, 1], [3, 1, 1, 1, 1]]>
user_142906810
Related