Function TEST

tensorflow/compiler/jit/shape_inference_test.cc:36–62

namespace {

TEST(ShapeInferenceTest, Basics) {
  Scope root = Scope::NewRootScope().ExitOnError();
  auto a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT,
                            ops::Placeholder::Shape({2, 3}));
  auto b = ops::Placeholder(root.WithOpName("B"), DT_FLOAT,
                            ops::Placeholder::Shape({3}));
  auto c = ops::Placeholder(root.WithOpName("C"), DT_FLOAT);
  auto d = ops::Add(root.WithOpName("D"), a, b);
  auto e = ops::Add(root.WithOpName("E"), d, c);
  auto f = ops::Neg(root.WithOpName("F"), e);
  auto g = ops::AddN(root.WithOpName("G"), std::initializer_list<Output>{e, f});

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

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(graph.get(), /*arg_shapes=*/{},
                           /*fnlib_def=*/nullptr, &shape_info));

  std::map<string, std::vector<PartialTensorShape>> expected = {
      {"A", {PartialTensorShape({2, 3})}}, {"B", {PartialTensorShape({3})}},
      {"C", {PartialTensorShape()}},       {"D", {PartialTensorShape({2, 3})}},
      {"E", {PartialTensorShape()}},       {"F", {PartialTensorShape()}},
      {"G", {PartialTensorShape()}},
  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(*graph, shape_info, expected));
}

TEST(ShapeInferenceTest, WhileLoop) {
  // Graph: