How to decorate a function that takes a tf.variable as a parameter with tf.function and most importantly using input signature

Question

I have a problem where I need to modify a variable inside a Tensorflow function. Then I need to convert this function to a tensorflow graph. The problem is that the size of the variable is not fix. Example: it can be either a tenosr of shape (3,) or (2,). This is why the function takes this variable as a parameter, so that it can modify it and return it.

Here is an example of a class that contains a function call, this function takes two arguments (x,v). x is a Tf.tensor and v is a tf.Variable. v is assigned the the multiplication of x*v.

import tensorflow as tf

class MyModule(tf.Module):
  def __init__(self):
    pass

  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.int32), tf.TensorSpec(shape=[None], dtype=tf.int32)])
  def __call__(self, x, v):
    v.assign(x*v, read_value=False)
    return v

tf.config.run_functions_eagerly(False)
x = tf.constant([10,10])
v = tf.Variable(2*tf.ones_like(x), trainable=False)

module = MyModule()
module(x, v)

This works as expected in eager mode, but in graph mode I get the following error: AttributeError: 'Tensor' object has no attribute 'assign'

I know that it is because of the signature of tf.Variable. My question is how can I specify the signature of tf.Variable given that the current one produces an error?

Answer 1

Actually there is one operation that can achieve what you want, however it is not listed in the public API. Beware that may not be the best practice.

You need resource_variable_ops which you can find under tensorflow.python.ops.

import tensorflow as tf
from tensorflow.python.ops import resource_variable_ops

class MyModule(tf.Module):
  def __init__(self):
    pass

  @tf.function(input_signature=[
                                tf.TensorSpec(shape=[None], dtype=tf.int32), 
                                resource_variable_ops.VariableSpec(shape=[None], dtype=tf.int32)
                                ])
  def __call__(self, x, v):
    v.assign(x*v, read_value=False)
    return v

x = tf.constant([10,10])
v = tf.Variable(2*tf.ones_like(x), trainable=False)

module = MyModule()
module(x, v)

Answer 2

Here is the solution I have found:

class MyModule(tf.Module):
    def __init__(self):
        self.v = tf.Variable([[]], shape=tf.TensorShape([None]*2), dtype=tf.int32)

    @tf.function(input_signature=[tf.TensorSpec(shape=[None]*2, dtype=tf.int32), tf.TensorSpec(shape=[None]*2, dtype=tf.int32)])
    def __call__(self, x, v):
      self.v.assign( x * v, read_value=False)
      return self.v

Defining tf.Variable inside the constructor of the model with empty values solved the problem.

Answer 3

you can complete it with a simple method.

[ Sample ]:

import tensorflow as tf

class MyModule(tf.Module):
    def __init__(self, v):
        self.v = v
        pass

    @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.int32), tf.TensorSpec(shape=None, dtype=tf.int32)])
    def __call__(self, x, v):
        self.v.assign( x * v, read_value=False )
        return self.v

x = tf.constant( tf.random.uniform(shape=[2,1], maxval=3, dtype=tf.int32) )
v = tf.Variable([[1], [2]])
module = MyModule(v)

print( module(x, v) )

#############################################
x = tf.constant( tf.random.uniform(shape=[3,1], maxval=3, dtype=tf.int32) )
v = tf.Variable([[1], [2], [3]])
module = MyModule(v)

print( module(x, v) )

[ Output ]:

tf.Tensor(
[[1]
 [0]], shape=(2, 1), dtype=int32)
tf.Tensor(
[[2]
 [0]
 [6]], shape=(3, 1), dtype=int32)