tensorflow: convert a list of tensor to ragged tensor with a fixed dim in a certain axis-CodePudding

Suppose I have three numpy arrays with shape (1,3) and stack them into groups with shape (2,3) and (1,3). Then I stack them with tf.ragged.stack to get a ragged tensor:

x1 = np.asarray([1,0,0])
x2 = np.asarray([0,1,0])
x3 = np.asarray([0,0,1])

group_a = np.stack([x1,x2])
group_b = np.stack([x3])


ac = tf.ragged.stack([group_a,group_b], axis=0)

I expect its shape to be (2, None, 3) but instead it's (2, None, None). How do I get the desired shape? I'm using tensorflow 2.5.2

CodePudding user response：

This is happening because tf.ragged.stack is creating a ragged_rank which equals to 2. Check the docs for more information. You could explicitly define how to partition a ragged tensor like this:

import tensorflow as tf
import numpy as np

x1 = np.asarray([1,0,0])
x2 = np.asarray([0,1,0])
x3 = np.asarray([0,0,1])

ac = tf.RaggedTensor.from_row_splits(
    values=[x1, x2, x3],
    row_splits=[0, 2, 3])

print(ac.shape)
print(ac)

(2, None, 3)
<tf.RaggedTensor [[[1, 0, 0], [0, 1, 0]], [[0, 0, 1]]]>