MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / AddSwitchNodeAlongEdge

Method AddSwitchNodeAlongEdge

tensorflow/compiler/tf2xla/functionalize_cond.cc:484–511  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

482}
483
484Status Conditional::AddSwitchNodeAlongEdge(const Edge* edge, BranchType branch,
485 Graph* graph) {
486 // Previously we had edge:
487 // src:src_output ---- edge ----> dst:dst_input
488 // post this we have (in graph)
489 // src:src_output --> switch<pred> --- new_edge --> dst:dst_input
490
491 // TODO(jpienaar): One could keep a map caching the extra switch nodes added
492 // to avoid adding another switch to feed a value for which a switch was
493 // already added.
494 Node* switch_node;
495 Node* src = edge->src();
496 int src_output = edge->src_output();
497 TF_RETURN_IF_ERROR(
498 NodeBuilder(graph->NewName(absl::StrCat(src->name(), "_added_switch")),
499 "Switch")
500 .Input(src, src_output)
501 .Input(const_cast<Node*>(predicate_.node), predicate_.index)
502 .Finalize(graph, &switch_node));
503 state_map_->ResetCondId(switch_node, state_map_->LookupCondId(src));
504 state_map_->ResetAncestorId(switch_node, state_map_->LookupAncestorId(src));
505
506 Node* dst = edge->dst();
507 int dst_input = edge->dst_input();
508 graph->RemoveEdge(edge);
509 graph->AddEdge(switch_node, static_cast<int>(branch), dst, dst_input);
510 return AddSwitch(switch_node);
511}
512
513Status Conditional::ExtractBodies(Graph* graph) {
514 VLOG(2) << "Extracting bodies for " << name();

Callers

nothing calls this directly

Calls 15

ResetCondIdMethod · 0.80
LookupCondIdMethod · 0.80
ResetAncestorIdMethod · 0.80
LookupAncestorIdMethod · 0.80
nameMethod · 0.65
NodeBuilderClass · 0.50
StrCatFunction · 0.50
srcMethod · 0.45
src_outputMethod · 0.45
FinalizeMethod · 0.45
InputMethod · 0.45
NewNameMethod · 0.45

Tested by

no test coverage detected