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

Question

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

Answer 1

This question is actually answered by the official help: https://www.tensorflow.org/api_docs/python/tf/RaggedTensor#get_shape

x1=[[0, 1, 2]]
x2=[[1, 2, 3], [3, 4, 6]]

tf.ragged.constant([x1, x2], ragged_rank=1).shape

Out[51]: TensorShape([2, None, 3])