Let Defun helper methods able to recognize TF's new FuncGraph.

PiperOrigin-RevId: 309148611
Change-Id: I4ea99d7b6ac77352fc32422ea60f5de795075e62
This commit is contained in:
A. Unique TensorFlower 2020-04-29 19:57:57 -07:00 committed by TensorFlower Gardener
parent 2d80f55c77
commit 8314abe2e5

View File

@ -31,7 +31,6 @@ from tensorflow.python.client import pywrap_tf_session as c_api
from tensorflow.python.eager import context
from tensorflow.python.framework import c_api_util
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import func_graph as _func_graph
from tensorflow.python.framework import graph_to_function_def
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
@ -1274,8 +1273,6 @@ def get_extra_vars():
g = ops.get_default_graph()
if isinstance(g, _FuncGraph):
return g.extra_vars
elif isinstance(g, _func_graph.FuncGraph):
return g.variable_captures
else:
return []
@ -1292,8 +1289,6 @@ def get_extra_inputs():
g = ops.get_default_graph()
if isinstance(g, _FuncGraph):
return g.extra_inputs
elif isinstance(g, _func_graph.FuncGraph):
return g.external_captures
else:
return []
@ -1310,8 +1305,6 @@ def get_extra_args():
g = ops.get_default_graph()
if isinstance(g, _FuncGraph):
return g.extra_args
elif isinstance(g, _func_graph.FuncGraph):
return g.internal_captures
else:
return []