有什么漂亮的方法可以让 tf.function 接受类实例值吗？

Question

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

有什么办法让这成为可能吗？目前，我只是将add_constant内部装饰Foo class __init__为：

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

然而，与函数装饰器相比，这使得代码的可读性较差。