Converts a shape inference handle to a PartialTensorShape.
| 27 | |
| 28 | // Converts a shape inference handle to a PartialTensorShape. |
| 29 | Status ShapeHandleToTensorShape(shape_inference::InferenceContext* context, |
| 30 | const shape_inference::ShapeHandle& handle, |
| 31 | PartialTensorShape* shape) { |
| 32 | // The default is already unknown |
| 33 | if (!context->RankKnown(handle)) return Status::OK(); |
| 34 | |
| 35 | std::vector<int64> dims(context->Rank(handle)); |
| 36 | for (int32 i = 0; i < dims.size(); ++i) { |
| 37 | dims[i] = context->Value(context->Dim(handle, i)); |
| 38 | } |
| 39 | return PartialTensorShape::MakePartialShape(dims.data(), dims.size(), shape); |
| 40 | } |
| 41 | |
| 42 | Status PropagateShapes(const Graph& graph, |
| 43 | const std::map<int, InferredShape>& arg_shapes, |