Any beautiful way to make tf.function takes in class instance value?

Viewed 1353
class Foo(object):
    def __init__(self):
        self.shape = [None, 4]

    @tf.function(input_signature=[tf.TensorSpec(shape=self.shape)])
    def add_constant(self, a):
        print('Building graph')
        return a + 5

Any way to make this possible? Currently, I just decorate add_constant inside Foo class __init__ as:

self.add_constant = tf.function(func=self.add_constant, input_signature=[tf.TensorSpec(shape=self.shape)]).

However, this makes the code less readable comparable to function decorator.

0 Answers
Related