Expose @tf.function annotated methods on tf.Module as functions.
An issue with decorating methods in the metaclass is that when users annotate methods with decorators that return Objects (e.g. `tf.function`) they lose the ability to access methods on those objects (other than call) because our `with_name_scope` decorator is "on the outside". A simple answer to this from other decorators is to copy the __dict__ from the wrapped object to the wrapper (or use `__getattr__`). However in our case this is not safe, since properties like `get_concrete_function` on `Function` would need special casing (so the returned callable also entered the module name scope before tracing). After a bunch of discussion (thanks superbobry@ and chr1sj0nes@!) about how we could do this in horrible ways, it seems that the cleanest solution is to allow a Function to be rewrapped (iff it has not been traced). This pushes the decorator down the stack of decorators and allows tf.Module to preserve @tf.function annoated methods as Functions. h/t tensorflow/community#56 PiperOrigin-RevId: 233034181
Loading
Please sign in to comment