Skip to content
Permalink
Browse files

Fix a bug in tpu.py and xla.py that while creating an identity node f…

…or control input edges under rewrite context, the parent control flow context is lost. (#23446)

PiperOrigin-RevId: 219724472
  • Loading branch information...
tatatodd committed Nov 2, 2018
1 parent 8ce231a commit a6d8ffae097d0132989ae4688d224121ec6d8f35
Showing with 10 additions and 16 deletions.
  1. +5 −8 tensorflow/contrib/compiler/xla.py
  2. +5 −8 tensorflow/contrib/tpu/python/tpu/tpu.py
@@ -179,14 +179,11 @@ def AddOp(self, op):
if external_control_inputs:
# Use an identity to pull control inputs as data inputs. Note that we
# ignore ops which don't have outputs. TODO(phawkins): fix that.
with ops.control_dependencies(None):
self.Enter()
external_control_inputs = [
array_ops.identity(x.outputs[0]).op
for x in external_control_inputs
if x.outputs
]
self.Exit()
external_control_inputs = [
array_ops.identity(x.outputs[0]).op
for x in external_control_inputs
if x.outputs
]
# pylint: disable=protected-access
op._add_control_inputs(external_control_inputs)
# pylint: enable=protected-access
@@ -371,14 +371,11 @@ def AddOp(self, op):
if external_control_inputs:
# Use an identity to pull control inputs as data inputs. Note that we
# ignore ops which don't have outputs. TODO(phawkins): fix that.
with ops.control_dependencies(None):
self.Enter()
external_control_inputs = [
array_ops.identity(x.outputs[0]).op
for x in external_control_inputs
if x.outputs
]
self.Exit()
external_control_inputs = [
array_ops.identity(x.outputs[0]).op
for x in external_control_inputs
if x.outputs
]
# pylint: disable=protected-access
op._add_control_inputs(external_control_inputs)
# pylint: enable=protected-access

0 comments on commit a6d8ffa

Please sign in to comment.
You can’t perform that action at this time.