Define `input_signature` with multiple input shapes

General Discussion
#1

Is it possible to define the input_signature of a @tf.function which the input tensor can have three different shapes (e.g., [B, 2], [B, H, W, 33], and [B, H, W, 2]?

I want to define this in the call method:

@tf.function(input_signature = [ tf.TensorSpec(shape=[], dtype=tf.float32, name="input") ])
def call(self, tensor: tf.Tensor) -> tf.Tensor:
    ...
#3

Hi @fabricionarcizo, you can define the input_signature of a tf.function with different shapes. For example,

@tf.function(input_signature=(tf.TensorSpec(shape=[1,2], dtype=tf.float32),
                              tf.TensorSpec(shape=[1,128,128,33], dtype=tf.float32),
                              tf.TensorSpec(shape=[1,128,128,2], dtype=tf.float32)))
def my_function(input_tensor1,input_tensor2,input_tensor3):
    return input_tensor1

Thank You.