pythontensorflowragged

Unnest ragged tensor in tensorflow

I have following tf.RaggedTensor rt:

<tf.RaggedTensor [[[7592]], [[7592], [2088]], [[7592], [2045], [2236]], [[7592], [2045], [2236], [6358, 16429, 2072]]]>

here is the dense version:

array([[[ 7592,     0,     0],
        [    0,     0,     0],
        [    0,     0,     0],
        [    0,     0,     0]],

       [[ 7592,     0,     0],
        [ 2088,     0,     0],
        [    0,     0,     0],
        [    0,     0,     0]],

       [[ 7592,     0,     0],
        [ 2045,     0,     0],
        [ 2236,     0,     0],
        [    0,     0,     0]],

       [[ 7592,     0,     0],
        [ 2045,     0,     0],
        [ 2236,     0,     0],
        [ 6358, 16429,  2072]]], dtype=int32)>

I need to "unnest" rt to this:

<tf.RaggedTensor [[7592], [7592,2088], [7592,2045,2236], [7592,2045,2236,6358, 16429, 2072]]>

but I've had no luck with combinations of tf.concat, tf.ragged.map_flat_values or tf.squeeze.

Anyone know the answer?

Solution

  • After some searching, turns out there's a built-in method: rt.merge_dims(1,2).