[SOLVED] Tensorflow concat ragged tensor returning wrong shape

Tensorflow concat ragged tensor returning wrong shape

I need to concat two ragged tensors keeping the last dimension with a fixed size 2.
Checking model.output_shape I get the desired (None, None, 2). But when I call the model, I get (batch_size, None, None). How do I get the right shape? Code:

import tensorflow as tf

a_input = tf.keras.layers.Input([None, 2], ragged=True)
b_input = tf.keras.layers.Input([None, 2], ragged=True)
output = tf.concat([a_input, b_input], axis=1)

model = tf.keras.Model([a_input, b_input], output)

a = tf.ragged.constant([
    [[1, 2], [3, 4], [5, 6]],
    [[1, 2], [3, 4]],
    [[1, 2]],
])
b = tf.ragged.constant([
    [[1, 2]],
    [[1, 2], [3, 4], [5, 6], [7, 8]],
    [[1, 2], [3, 4]],
])

print(model.output_shape)
# (None, None, 2)
print(model([a, b]).shape)
# (3, None, None)

Solution

I found it. The tf.ragged.constant does not consider the last dimension a uniform dimension. So a.shape is (3, None, None). To fix that I need to use ragged_rank parameter:

a = tf.ragged.constant([
    [[1, 2], [3, 4], [5, 6]],
    [[1, 2], [3, 4]],
    [[1, 2]],
], ragged_rank=1)

print(a.shape)
# (3, None, 2)