[py_function] Don't attach py_function to the global eager graph.
Eager mode can incorrectly have a global graph. Disabling global graph
on eager mode breaks too many assumptions so first introduce a flag indicating it.
Also, avoid attaching py_function to eager mode global graph, which is a leak.
Though this CL doesn't fix the leak yet as there are two more references that leads
to the leak, `tape_cache` and `ag_dnc_wrapper__` .
#35084
PiperOrigin-RevId: 288728035
Change-Id: I27c254de4323e3fcac9966294e624dda61f91cd2
diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index 1d77e71..f50ffa0 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -2788,11 +2788,6 @@
# tuples: (input_shape_tuple, reduction_indices_tuple), and the values
# are pairs of tuples: (output_shape_kept_dims, tile_scaling).
self._reduced_shape_cache = {}
- # In eager mode, the top level graph can still be created. This is
- # incorrect and undesriable but currently so many places are relying on
- # this. This is a flag indicating that, and meant to be set manually after
- # this graph construction.
- self._is_eager_graph = False
# TODO(skyewm): fold as much of the above as possible into the C
# implementation
@@ -5364,8 +5359,6 @@
# the global default graph and an explicit graph are combined in the
# same process.
self._global_default_graph = Graph()
- if context.executing_eagerly():
- self._global_default_graph._is_eager_graph = True # pylint: disable=protected-access
return self._global_default_graph
def reset(self):
diff --git a/tensorflow/python/ops/script_ops.py b/tensorflow/python/ops/script_ops.py
index 16711e6..8463ffb 100644
--- a/tensorflow/python/ops/script_ops.py
+++ b/tensorflow/python/ops/script_ops.py
@@ -316,11 +316,9 @@
while True:
current_graph = graph
if isinstance(graph, function._FuncGraph): # pylint: disable=protected-access
- if not graph._outer_graph._is_eager_graph: # pylint: disable=protected-access
- graph = graph._outer_graph # pylint: disable=protected-access
+ graph = graph._outer_graph # pylint: disable=protected-access
elif isinstance(graph, func_graph.FuncGraph):
- if not graph.outer_graph._is_eager_graph: # pylint: disable=protected-access
- graph = graph.outer_graph
+ graph = graph.outer_graph
if graph is current_graph:
break
