Commit 810a46c0 authored by Tom Hennigan's avatar Tom Hennigan Committed by TensorFlower Gardener
Browse files

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
parent 4f1ca2f2
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment