| 156 | } |
| 157 | |
| 158 | Status CondBuilder::AddInput(Node* src, int src_output) { |
| 159 | Node* input; |
| 160 | NodeDebugInfo debug_info(*src); |
| 161 | // Colocate the Switch node with the `src` node. |
| 162 | // |
| 163 | // This is to avoid unnecessary Host<->Device copies between src and the |
| 164 | // Switch node. This aligns with the implementation of legacy tf.cond in |
| 165 | // control_flow_ops.py. The legacy impl colocates the Switch with the |
| 166 | // input tensor which resets the device stack and forces the Switch to have |
| 167 | // the same device as the input node (if set) and sets the colocation _class |
| 168 | // attr. It also ignores the existing colocation constraints on the input node |
| 169 | // using colocate_with(ignore_existing=True). |
| 170 | TF_RETURN_IF_ERROR(NodeBuilder(NewName(src->name()), "Switch", |
| 171 | graph_->op_registry(), &debug_info) |
| 172 | .Input(src, src_output) |
| 173 | .Input(pred_) |
| 174 | .Device(src->requested_device()) |
| 175 | .Attr("_class", {src->name()}) |
| 176 | .Finalize(graph_, &input)); |
| 177 | then_call_builder_.Input(input, kThenBranch); |
| 178 | else_call_builder_.Input(input, kElseBranch); |
| 179 | return Status::OK(); |
| 180 | } |
| 181 | |
| 182 | Status CondBuilder::AddInputs() { |
| 183 | // Add input data edges. |
nothing calls this directly
no test coverage detected