More thoroughly disable the should_use_result decorator when executing eagerly.

It was creating reference cycles.

Adds a test that TensorArrays create no reference cycles in eager mode.

PiperOrigin-RevId: 174768765

1 files changed, 11 insertions, 3 deletions

diff --git a/tensorflow/python/util/tf_should_use.py b/tensorflow/python/util/tf_should_use.py
index 99081cb294..a576547d5f 100644
--- a/tensorflow/python/util/tf_should_use.py
+++ b/tensorflow/python/util/tf_should_use.py
@@ -22,6 +22,7 @@ import types
 
 import six  # pylint: disable=unused-import
 
+from tensorflow.python.eager import context
 from tensorflow.python.util import tf_decorator
 # pylint: enable=g-bad-import-order,g-import-not-at-top
 
@@ -31,6 +32,8 @@ from tensorflow.python.util import tf_decorator
 def _add_should_use_warning(x, fatal_error=False):
   """Wraps object x so that if it is never used, a warning is logged.
 
+  Does nothing when executing eagerly.
+
   Args:
     x: Python object.
     fatal_error: Python bool.  If `True`, tf.logging.fatal is raised
@@ -44,9 +47,10 @@ def _add_should_use_warning(x, fatal_error=False):
   if x is None:  # special corner case where x is None
     return x
 
-  # TODO(apassos) we don't have an easier way to check because importing context
-  # or ops here would create a BUILD dependency cycle.
-  if type(x).__name__ == 'EagerTensor':
+  if context.in_eager_mode():
+    # Typically not needed when executing eagerly (the main use case is for ops
+    # which need to be incorporated into the graph), and even the no-op wrapper
+    # creates reference cycles which require garbage collection.
     return x
 
   def override_method(method):
@@ -102,6 +106,8 @@ def should_use_result(fn):
   - `t != 0`.  In this case, comparison is done on types / ids.
   - `isinstance(t, tf.Tensor)`.  Similar to above.
 
+  Does nothing when executing eagerly.
+
   Args:
     fn: The function to wrap.
 
@@ -136,6 +142,8 @@ def must_use_result_or_fatal(fn):
   - `t != 0`.  In this case, comparison is done on types / ids.
   - `isinstance(t, tf.Tensor)`.  Similar to above.
 
+  Does nothing when executing eagerly.
+
   Args:
     fn: The function to wrap.