{"C", {PartialTensorShape()}},       {"D", {PartialTensorShape({2, 3})}},
      {"E", {PartialTensorShape()}},       {"F", {PartialTensorShape()}},
      {"G", {PartialTensorShape()}},
  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(*graph, shape_info, expected));
}

// Test that shape inference uses user-given `arg_shapes` to propagate shapes.
TEST(ShapeInferenceTest, UseArgShapesForVariableBatchSize) {

// `expected_shapes`. Returns an error if there are nodes in `expected_shapes`
// that do not have shape information. Ignores nodes in `graph` that do not have
// `expected_shapes` entries.
Status ShapeAnnotationsMatch(
    const Graph& graph, const GraphShapeInfo& shape_info,
    std::map<string, std::vector<PartialTensorShape>> expected_shapes);

// A helper object to create GraphOptimizationPassOptions.

#include "tensorflow/compiler/jit/shape_inference.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/public/version.h"

namespace tensorflow {

Status ShapeAnnotationsMatch(
    const Graph& graph, const GraphShapeInfo& shape_info,
    std::map<string, std::vector<PartialTensorShape>> expected_shapes) {
  for (Node* node : graph.op_nodes()) {
    auto sit = shape_info.find(node->name());