Function TEST_F

tensorflow/core/grappler/optimizers/shape_optimizer_test.cc:33–73

class ShapeOptimizerTest : public GrapplerTest {};

TEST_F(ShapeOptimizerTest, OptimizeShapeProduct) {
  // Set the device to CPU zero, because the shape optimizer will only optimize
  // Prod to Size when a concrete Size kernel is available.
  tensorflow::Scope s = tensorflow::Scope::NewRootScope().WithDevice("/cpu:0");
  Output a = ops::Const(s.WithOpName("a"), 3.14f, {32, 16});
  Output c = ops::Shape(s.WithOpName("c"), a);
  Output d = ops::Const(s.WithOpName("d"), 0, {1});
  ops::ReduceProd::Attrs attrs;
  Output e = ops::ReduceProd(s.WithOpName("e"), c, d, attrs.KeepDims(false));
  Output f = ops::ReduceProd(s.WithOpName("f"), c, d, attrs.KeepDims(true));

  GrapplerItem item;
  item.fetch = {"e", "f"};
  TF_CHECK_OK(s.ToGraphDef(&item.graph));

  auto tensors_expected = EvaluateNodes(item.graph, item.fetch);

  GraphDef output;
  ShapeOptimizer optimizer;
  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));

  int found = 0;
  for (const NodeDef& node : output.node()) {
    if (node.name() == "e") {
      found++;
      EXPECT_EQ("Size", node.op());
      EXPECT_EQ("a", node.input(0));
    } else if (node.name() == "f") {
      found++;
      EXPECT_EQ("Prod", node.op());
      EXPECT_EQ("c", node.input(0));
    }
  }
  EXPECT_EQ(2, found);

  auto tensors_actual = EvaluateNodes(output, item.fetch);
  EXPECT_NEAR(tensors_expected[0].scalar<int>()(),
              tensors_actual[0].scalar<int>()(), 0);
  EXPECT_NEAR(tensors_expected[1].scalar<int>()(),
              tensors_actual[1].scalar<int>()(), 0);
}

TEST_F(ShapeOptimizerTest, OptimizeShapeProductMissingKernel) {
  {