diff options
| author |  Skye Wanderman-Milne <skyewm@google.com> | 2017-12-08 15:45:25 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-12-08 15:49:02 -0800 |
| commit | 28807c5666c9f574ef415fed7b18b99ebed41ecc (patch) | |
| tree | 15c7249d2925465202e21ac800a81a6df5bda8f2 /tensorflow/c/python_api.cc | |
| parent | b1c7d177e2aa9a4e3989caf7cfb21a5591c3832f (diff) | |
Add Operation._remove_all_control_inputs and use in ControlFlowContext.
This allows while loop gradients to work with the C API. This change
also enables the C API for control flow tests.
PiperOrigin-RevId: 178438424
Diffstat (limited to 'tensorflow/c/python_api.cc')
| -rw-r--r-- | tensorflow/c/python_api.cc | 12 |
1 files changed, 12 insertions, 0 deletions
diff --git a/tensorflow/c/python_api.cc b/tensorflow/c/python_api.cc index 37629a74ba..6e37cdb5f4 100644 --- a/tensorflow/c/python_api.cc +++ b/tensorflow/c/python_api.cc @@ -87,4 +87,16 @@ void UpdateEdge(TF_Graph* graph, TF_Output new_src, TF_Input dst, } } +void RemoveAllControlInputs(TF_Graph* graph, TF_Operation* op) { + mutex_lock l(graph->mu); + std::vector<const Edge*> control_edges; + for (const Edge* edge : op->node.in_edges()) { + if (!edge->IsControlEdge()) continue; + control_edges.push_back(edge); + } + for (const Edge* edge : control_edges) { + graph->graph.RemoveControlEdge(edge); + } +} + } // namespace tensorflow |