Naming and sharing unknown dimension across multiple inputs.

[mkschleg]

Ask a Question

Question

Is there a way to label an unknown dimensions as the same across a model? For instance, if there is a batchsize that is unknown when converting the model but shared across many inputs.

Further information

For example:

import tensorflow as tf
import tf2onnx

@tf.function
def my_func(x, y):
    i = x.shape[0]
    j = y.shape[0]
    ret_1 = tf.reduce_sum(x, axis=0)
    ret_2 = tf.pow(y, 2)
    return i, j, ret_1, ret_2

tf2onnx.convert.from_function(
    my_func, opset=13, input_signature=[tf.TensorSpec((None, 2), tf.float32), tf.TensorSpec((None,), tf.int32)], output_path="test.onnx")

Both x and y in the above example have a shape where the first dimension is unknown (in my converted model they are named unk__13 and unk__14 respectively). Is there a way to let tf2onnx know that these should actually be the same size (i.e. unk__13 == unk__14).

Notes

I'm not sure if this effects performance, but would make interpreting the resulting onnx models easier for downstream users.