How would I make a "dynamic" function signature in TensorFlow?

user_10935831 0
Viewed 95

I've been trying to make a "dynamic" function signature with TensorFlow using @tf.function(input_signature=[...]). More specifically, I need to find a way for one of the arguments to be either tf.TensorSpec(shape=[None, None], dtype=tf.float32) or None

I am currently trying the following:

@tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec(shape=[None, None], dtype=tf.float32)])
def generate_one_step(self, inputs, states=None):
  ...

but as the second parameter could also in some cases be None I get the following error:

ValueError: When input_signature is provided, all inputs to the Python function must be convertible to tensors:
  inputs: (
    tf.Tensor([b'Test '], shape=(1,), dtype=string),
    None)
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.string, name=None),
    TensorSpec(shape=(None, None), dtype=tf.float32, name=None)).

I realize there may also be an error regarding the shape of the first input but I am unsure how to make a "variable" shape as well

tensorflow python
1 Answers

The workaround for this turned out to be to just specify two different functions. It is worth pointing out that TF signatures do NOT work like Java signatures, so you should name your functions differently.

My code:

  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec(shape=[None, None], dtype=tf.float32)])
  def generate_one_step(self, inputs, states):

  ...
  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)])
  def generate_one_step_none(self, inputs):
  states = None
user_109358310
Related