tf.function with a dynamic number of dynamic-shaped tensors

Question

I have some code that runs fine in eager mode, which processes a list of tensors. The length of the list is known at runtime only, and the list elements have different shapes (e.g. different row number but same number of cols), also known at runtime. I would like to write a graph version of this code.

To avoid re-tracing at each call with different list length / element dimensions I would like to use an input_signature with None at the dynamic dimensions. A list of tensors is not an acceptable signature, so I tried two different options:

Using a RaggedTensor, and a RaggedTensorSpec in the input_signature, but the problem is that None in that case already means a ragged axis, and there is no value I can use for dimensions that are not ragged but not known yet.
Using a TensorArray (with infer_shape=False), but I get an error when using a TensorArraySpec in the input_signature and calling the function (maybe it is not possible to use a TensorArraySpec here at all ?).

Here is a minimal example:

a = tf.constant([[1, 2, 3],
                 [4, 5, 6]], dtype=tf.float32) #Shape [2, 3]

b = tf.constant([[7, 8, 9]], dtype=tf.float32) #Shape [1, 3]


#Version with ragged tensor

values = tf.concat([a, b], axis=0)
lens = tf.stack([tf.shape(t)[0] for t in [a, b]])
ragged_tensor = tf.RaggedTensor.from_row_lengths(values, lens)
ragged_tensor.shape #TensorShape([2, None, 3])
#ragged_tensor represents a list of 2 tensors of shape [None, 3] (the 2 can vary but the 3 is fixed)

@tf.function(input_signature=(tf.RaggedTensorSpec(shape=[None, None, 3], dtype=tf.float32),))
def myFunc(x):
    return tf.math.reduce_sum(x)

myFunc(ragged_tensor) #ValueError


#Version with tensor array

array = tf.TensorArray(tf.float32, size=2, infer_shape=False, element_shape=[None, 3])
array = array.write(0, a)
array = array.write(1, b)
#array represents a list of 2 tensors of shape [None, 3] (the 2 can vary but the 3 is fixed)

@tf.function(input_signature=(tf.TensorArraySpec(element_shape=[None, 3], dtype=tf.float32, infer_shape=False),))
def myFunc1(x):
    return tf.math.reduce_sum(x.concat())

myFunc1(array) #TypeError

The full error messages:

ValueError: Incompatible CompositeTensor TypeSpecs: RaggedTensorSpec(TensorShape([None, None, 3]), tf.float32, 2, tf.int64) vs. RaggedTensorSpec(TensorShape([2, None, 3]), tf.float32, 1, tf.int32) -- Types are not compatible: 2 vs 1

TypeError:If shallow structure is a sequence, input must also be a sequence. Input has type:

How can I make it work ?

tf.function with a dynamic number of dynamic-shaped tensors

0 Answers