Drop external control dependencies in tfe.defun.

They shouldn't help given the automatic control dependencies, and are tricky
to capture in the general case.

PiperOrigin-RevId: 215282837

1 files changed, 9 insertions, 0 deletions

diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 3b6f288fb9..f261d92d64 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -269,6 +269,15 @@ class FuncGraph(ops.Graph):
   def variables(self, var_list):
     self._weak_variables = [weakref.ref(v) for v in var_list]
 
+  def control_dependencies(self, control_inputs):
+    # Drop control dependencies to outside of the graph. TODO(b/117109273)
+    # unclear how to capture an op, not a tensor.
+    if not control_inputs:
+      return super(FuncGraph, self).control_dependencies(control_inputs)
+    return super(FuncGraph, self).control_dependencies(
+        [c for c in control_inputs
+         if getattr(c, "graph", None) is self])
+
   def create_op(
       self,
       op_type,