| 145 | } |
| 146 | |
| 147 | Status CaseBuilder::AddInput(Node* src, int src_output) { |
| 148 | Node* input; |
| 149 | NodeDebugInfo debug_info(*src); |
| 150 | // Colocate the Switch node with the `src` node. |
| 151 | // |
| 152 | // This is to avoid unnecessary Host<->Device copies between src and the |
| 153 | // _SwitchN node. This aligns with the implementation of legacy tf.cond in |
| 154 | // control_flow_ops.py. The legacy impl colocates the Switch with the |
| 155 | // input tensor which resets the device stack and forces the Switch to have |
| 156 | // the same device as the input node (if set) and sets the colocation _class |
| 157 | // attr. It also ignores the existing colocation constraints on the input node |
| 158 | // using colocate_with(ignore_existing=True). |
| 159 | TF_RETURN_IF_ERROR(NodeBuilder(NewName(src->name()), "_SwitchN", |
| 160 | graph_->op_registry(), &debug_info) |
| 161 | .Input(src, src_output) |
| 162 | .Input(branch_index_) |
| 163 | .Device(src->requested_device()) |
| 164 | .Attr("_class", {src->name()}) |
| 165 | .Attr("num_outs", num_branches_) |
| 166 | .Finalize(graph_, &input)); |
| 167 | for (int b = 0; b < num_branches_; b++) { |
| 168 | branch_call_builders_[b].Input(input, b); |
| 169 | } |
| 170 | return Status::OK(); |
| 171 | } |
| 172 | |
| 173 | Status CaseBuilder::AddInputs() { |
| 174 | // Add input data edges. |
nothing calls this directly
no test coverage detected