experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Handle the case when tf.function's target is a tf_decorator method that was hot-swapped after tf.function was applied.

Browse Source 
PiperOrigin-RevId: 235124593
This commit is contained in:
Dan Moldovan 2019-02-21 20:38:23 -08:00 committed by TensorFlower Gardener
parent 18f10890c0
commit 5295cfc096
1 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tensorflow/python/util/tf_decorator.py
View File

@ -59,6 +59,7 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import inspect
import traceback as _traceback
@ -154,7 +155,18 @@ def rewrap(decorator_func, previous_target, new_target):
return
target.decorated_target = new_target
innermost_decorator.__wrapped__ = new_target
if inspect.ismethod(innermost_decorator):
# Bound methods can't be assigned attributes. Thankfully, they seem to
# be just proxies for their unbound counterpart, and we can modify that.
if hasattr(innermost_decorator, '__func__'):
innermost_decorator.__func__.__wrapped__ = new_target
elif hasattr(innermost_decorator, 'im_func'):
innermost_decorator.im_func.__wrapped__ = new_target
else:
innermost_decorator.__wrapped__ = new_target
else:
innermost_decorator.__wrapped__ = new_target
def unwrap(maybe_tf_decorator):