| 2069 | } |
| 2070 | |
| 2071 | void MklLayoutRewritePass::GetNodesProducingTFTensorList( |
| 2072 | const gtl::InlinedVector<std::pair<Node*, int>, 4>& inputs, int* input_idx, |
| 2073 | int list_length, std::vector<NodeBuilder::NodeOut>* output_nodes) { |
| 2074 | CHECK_LT(*input_idx, inputs.size()); |
| 2075 | CHECK_GT(list_length, 0); |
| 2076 | CHECK_NOTNULL(output_nodes); |
| 2077 | output_nodes->reserve(list_length); |
| 2078 | |
| 2079 | while (list_length != 0) { |
| 2080 | CHECK_GT(list_length, 0); |
| 2081 | CHECK_LT(*input_idx, inputs.size()); |
| 2082 | Node* n = inputs[*input_idx].first; |
| 2083 | int slot = inputs[*input_idx].second; |
| 2084 | // If input node 'n' is just producing a single tensor at |
| 2085 | // output slot 'slot' then we just add that single node. |
| 2086 | output_nodes->push_back(NodeBuilder::NodeOut(n, slot)); |
| 2087 | (*input_idx)++; |
| 2088 | list_length--; |
| 2089 | } |
| 2090 | } |
| 2091 | |
| 2092 | // TODO(nhasabni) We should move this to mkl_util.h. |
| 2093 | void MklLayoutRewritePass::GetDummyMklTensorNode(std::unique_ptr<Graph>* g, |