| 218 | } // namespace |
| 219 | |
| 220 | Status ArgFeedRewrite::AddNode(Graph* g, NodeBuilder::NodeOut feed_tensor, |
| 221 | Node** out_node) { |
| 222 | // NOTE(mrry): We must include the index as part of the node |
| 223 | // name, because _Arg is a "stateful" kernel and therefore |
| 224 | // its name must uniquely identify a kernel instance across all |
| 225 | // graphs in the same session. |
| 226 | TF_RETURN_IF_ERROR( |
| 227 | NodeBuilder(strings::StrCat("_arg_", feed_tensor.node->name(), "_", |
| 228 | feed_tensor.index, "_", arg_index_), |
| 229 | "_Arg") |
| 230 | .Attr("T", BaseType(feed_tensor.node->output_type(feed_tensor.index))) |
| 231 | .Attr("index", arg_index_) |
| 232 | .Finalize(g, out_node, /*consume=*/true)); |
| 233 | (*out_node)->set_assigned_device_name(device_info().name()); |
| 234 | return Status::OK(); |
| 235 | } |
| 236 | |
| 237 | Status RecvFeedRewrite::AddNode(Graph* g, NodeBuilder::NodeOut feed_tensor, |
| 238 | Node** out_node) { |
no test coverage detected