| 180 | } // namespace |
| 181 | |
| 182 | Status InferShapes(Graph* graph, const std::map<int, InferredShape>& arg_shapes, |
| 183 | const tensorflow::FunctionLibraryDefinition* fnlib_def, |
| 184 | GraphShapeInfo* shape_info) { |
| 185 | ShapeRefiner shape_refiner(graph->versions(), graph->op_registry()); |
| 186 | shape_refiner.set_require_shape_inference_fns(false); |
| 187 | // TODO(dlibenzi): Verify if it is worth trying to infer shaped within |
| 188 | // functions. Some functions can be called at multiple locations with |
| 189 | // difference shapes, which will trigger a shape inference based on the |
| 190 | // arguments passed at the first call. |
| 191 | // shape_refiner.set_function_library_for_shape_inference(fnlib_def); |
| 192 | |
| 193 | // ShapeRefiner requires that all inputs of a node are present when |
| 194 | // ShapeRefiner::AddNode is called. To get at least some shape information in |
| 195 | // loops, we temporarily remove loop backedges and add them back again after |
| 196 | // the shape inference is complete. |
| 197 | BackEdgeHelper back_edge; |
| 198 | TF_RETURN_IF_ERROR(back_edge.Remove(graph)); |
| 199 | TF_RETURN_IF_ERROR(PropagateShapes(*graph, arg_shapes, |
| 200 | back_edge.RemovedEdges(), &shape_refiner)); |
| 201 | TF_RETURN_IF_ERROR(back_edge.Replace()); |
| 202 | |
| 203 | // Currently information does not flow "backward" from consumers to producers |
| 204 | // in the shape inference, but we consume the shapes in a second pass in case |
| 205 | // backward information flow is added in the future. |
| 206 | return StoreOutputShapes(*graph, shape_refiner, shape_info); |
| 207 | } |
| 208 | |
| 209 | xla::StatusOr<InferredShape> MergeInferredShapes(const InferredShape& a, |
| 210 | const InferredShape& b) { |