Store the shapes of the output tensors in a map
| 143 | |
| 144 | // Store the shapes of the output tensors in a map |
| 145 | Status StoreOutputShapes(const Graph& graph, const ShapeRefiner& shape_refiner, |
| 146 | GraphShapeInfo* shape_info) { |
| 147 | for (const Node* node : graph.nodes()) { |
| 148 | shape_inference::InferenceContext* context = shape_refiner.GetContext(node); |
| 149 | if (!context) continue; |
| 150 | |
| 151 | auto& outputs = (*shape_info)[node->name()]; |
| 152 | outputs.resize(context->num_outputs()); |
| 153 | for (int i = 0; i < context->num_outputs(); ++i) { |
| 154 | auto& output = outputs[i]; |
| 155 | TF_RETURN_IF_ERROR( |
| 156 | ShapeHandleToTensorShape(context, context->output(i), &output.shape)); |
| 157 | |
| 158 | const auto* handle_shapes_and_types = |
| 159 | context->output_handle_shapes_and_types(i); |
| 160 | if (handle_shapes_and_types != nullptr) { |
| 161 | if (handle_shapes_and_types->size() == 1) { |
| 162 | TF_RETURN_IF_ERROR(ShapeHandleToTensorShape( |
| 163 | context, (*handle_shapes_and_types)[0].shape, |
| 164 | &output.handle_shape)); |
| 165 | output.handle_type = (*handle_shapes_and_types)[0].dtype; |
| 166 | } else { |
| 167 | // otherwise, it may be resource like a Queue, which can have |
| 168 | // multiple shapes and types represented by a single handle. |
| 169 | } |
| 170 | } |
| 171 | VLOG(4) << node->name() << " output " << i << " shape" |
| 172 | << output.shape.DebugString() << " handle_type " |
| 173 | << DataTypeString(output.handle_type) << " handle_shape " |
| 174 | << output.handle_shape.DebugString(); |
| 175 | } |
| 176 | } |
| 177 | return Status::OK(); |
| 178 | } |
| 179 | |
| 180 | } // namespace |
| 181 |
no test coverage detected