Reshape Tensorflow RaggedTensor

Question

I have a 4D RaggedTensor of shape (batch_size, None, None, 100), I want to create from this a tensor of shape (batch_size, None, 100). So basically merging 1st and 2nd dimensions, but not including any padding ([1,2,3], [4] => [1,2,3,4]) and not converting to a dense tensor first. Is there a way to do this? If not what could be a work around?

Answer 1

After some more reading and trying, I have found the answer, which requires using row_starts twice for each of the two dimensions. The result loos like this:

row_starts = [my_ragged_tensor.values.row_starts()[row_start]
                      for row_start in my_ragged_tensor.row_starts()]

my_ragged_tensor_flat = tf.RaggedTensor.from_row_starts(my_ragged_tensor.flat_values, row_starts)

This will change the shape of the "my_ragged_tensor" from (batch_size, None, None, 100) to (batch_size, None, 100) merging the two middle dimensions.

Edit:

much easier way my_ragged_tensor.merge_dims(1,2)