Status InferShapes(Graph* graph, const std::map<int, InferredShape>& arg_shapes,
                   const tensorflow::FunctionLibraryDefinition* fnlib_def,
                   GraphShapeInfo* shape_info);

// Merges two InferredShapes. Return an error if the two shapes cannot be
// merged.
absl::StatusOr<InferredShape> MergeInferredShapes(const InferredShape& a,
                                                  const InferredShape& b);

  // backward information flow is added in the future.
  return StoreOutputShapes(*graph, shape_refiner, shape_info);
}

absl::StatusOr<InferredShape> MergeInferredShapes(const InferredShape& a,
                                                  const InferredShape& b) {
  InferredShape result;
  TF_RETURN_IF_ERROR(a.shape.MergeWith(b.shape, &result.shape));

  if (a.handle_type == DT_INVALID) {

  a.node()->AddAttr("_index", 0);
  b.node()->AddAttr("_index", 1);

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  TF_CHECK_OK(root.ToGraph(graph.get()));

  std::map<int, InferredShape> arg_shapes;
  arg_shapes[0].shape = TensorShape({2, 3});
  arg_shapes[1].shape = TensorShape({2, 3});

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(graph.get(), arg_shapes,

const char kXlaIsPlaceholderForArg[] = "_xla_is_placeholder_for_arg";

Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g) {
  // Perform shape inference.
  std::map<int, InferredShape> arg_shapes;
  GraphShapeInfo shape_info;
  TF_RETURN_IF_ERROR(
      InferShapes(g, arg_shapes, /*fnlib_def=*/nullptr, &shape_info));

  // Add attribute for output shapes.