| 2416 | } |
| 2417 | |
| 2418 | Status MklLayoutRewritePass::CopyInputs( |
| 2419 | const Node* old_node, |
| 2420 | const gtl::InlinedVector<std::pair<Node*, int>, 4>& old_node_inputs, |
| 2421 | NodeBuilder* nb) { |
| 2422 | // Number of input slots to old node |
| 2423 | // Input slots are represented by .Input() calls in REGISTER_OP. |
| 2424 | int old_node_input_slots = old_node->op_def().input_arg_size(); |
| 2425 | // Actual number of inputs can be greater than or equal to number |
| 2426 | // of Input slots because inputs of type list could be unfolded. |
| 2427 | auto old_node_input_size = old_node_inputs.size(); |
| 2428 | DCHECK_GE(old_node_input_size, old_node_input_slots); |
| 2429 | |
| 2430 | // Let's copy all inputs of old node to new node. |
| 2431 | int iidx = 0; |
| 2432 | for (int on_slot_idx = 0; on_slot_idx < old_node_input_slots; on_slot_idx++) { |
| 2433 | // An input slot could be a single tensor or a list. We need |
| 2434 | // to handle this case accordingly. |
| 2435 | DCHECK_LT(iidx, old_node_input_size); |
| 2436 | const OpDef::ArgDef& arg = old_node->op_def().input_arg(on_slot_idx); |
| 2437 | if (ArgIsList(arg)) { |
| 2438 | std::vector<NodeBuilder::NodeOut> new_node_inputs; |
| 2439 | int N = GetTensorListLength(arg, old_node); |
| 2440 | GetNodesProducingTFTensorList(old_node_inputs, &iidx, N, |
| 2441 | &new_node_inputs); |
| 2442 | nb->Input(new_node_inputs); |
| 2443 | } else { |
| 2444 | nb->Input(old_node_inputs[iidx].first, old_node_inputs[iidx].second); |
| 2445 | iidx++; |
| 2446 | } |
| 2447 | } |
| 2448 | return Status::OK(); |
| 2449 | } |
| 2450 | |
| 2451 | ////////////////////////////////////////////////////////////////////////// |
| 2452 | // Helper functions related to workspace pass |