Add EvaluateNodes to tests: AddOpsRewrite_AddOpsOfIdenticalShape, AddOpsRewrite_MultiplePasses, AddOpsRewrite_AddInputMultipleTimes, AddOpsRewrite_AddOpsOfSymbolicallyEqualShape, AddOpsRewrite_MinimizeBCast, AddOpsRewrite_MinimizeBCastWithSymbolicShapes, RemoveNegation, MinimizeBroadcasts_SimpleSwap, MinimizeBroadcasts_FlattenTallGraph, MinimizeBroadcasts_BuildTreeUp

PiperOrigin-RevId: 196125583

1 files changed, 130 insertions, 8 deletions

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
index 067adb359c..d60c3124ed 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
@@ -1574,6 +1574,14 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddOpsOfIdenticalShape) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto a_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto b_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto c_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  std::vector<std::pair<string, Tensor>> feed = {
+      {"a", a_t}, {"b", b_t}, {"c", c_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyAddToAddNCombining(&optimizer);
@@ -1607,6 +1615,10 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddOpsOfIdenticalShape) {
   ASSERT_NE(updated_outputs, nullptr);
 
   EXPECT_EQ(collapsed_add->name(), updated_outputs->input(0));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MultiplePasses) {
@@ -1631,6 +1643,17 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MultiplePasses) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto a_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto b_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto c_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto x_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto y_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto z_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  std::vector<std::pair<string, Tensor>> feed = {
+      {"a", a_t}, {"b", b_t}, {"c", c_t}, {"x", x_t}, {"y", y_t}, {"z", z_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyAddToAddNCombining(&optimizer);
@@ -1680,6 +1703,10 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MultiplePasses) {
   EXPECT_EQ(2, updated_mul->input_size());
   EXPECT_EQ(collapsed_left->name(), updated_mul->input(0));
   EXPECT_EQ(collapsed_right->name(), updated_mul->input(1));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddInputMultipleTimes) {
@@ -1697,6 +1724,14 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddInputMultipleTimes) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto a_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto b_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto c_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  std::vector<std::pair<string, Tensor>> feed = {
+      {"a", a_t}, {"b", b_t}, {"c", c_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyAddToAddNCombining(&optimizer);
@@ -1725,6 +1760,10 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddInputMultipleTimes) {
   EXPECT_EQ("b", collapsed_add->input(1));
   EXPECT_EQ("b", collapsed_add->input(2));
   EXPECT_EQ("c", collapsed_add->input(3));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddOpsOfSymbolicallyEqualShape) {
@@ -1748,6 +1787,11 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddOpsOfSymbolicallyEqualShape) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto x_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  std::vector<std::pair<string, Tensor>> feed = {{"input", x_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyAddToAddNCombining(&optimizer);
@@ -1779,6 +1823,10 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddOpsOfSymbolicallyEqualShape) {
   const NodeDef* updated_outputs = node_map.GetNode("outputs");
   ASSERT_NE(updated_outputs, nullptr);
   EXPECT_EQ(collapsed_add->name(), updated_outputs->input(0));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MinimizeBCast) {
@@ -1803,6 +1851,17 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MinimizeBCast) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto a_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32}));
+  auto b_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32, 32}));
+  auto c_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32, 32, 32}));
+  auto x_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32}));
+  auto y_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32, 32}));
+  auto z_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32, 32, 32}));
+  std::vector<std::pair<string, Tensor>> feed = {
+      {"a", a_t}, {"b", b_t}, {"c", c_t}, {"x", x_t}, {"y", y_t}, {"z", z_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyAddToAddNCombining(&optimizer);
@@ -1875,18 +1934,22 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MinimizeBCast) {
   const NodeDef* updated_outputs = node_map.GetNode("outputs");
   ASSERT_NE(updated_outputs, nullptr);
   EXPECT_EQ(outer_add_name, updated_outputs->input(0));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MinimizeBCastWithSymbolicShapes) {
   tensorflow::Scope s = tensorflow::Scope::NewRootScope();
 
   // We have a small input with one unknown dimension
-  auto small = ops::Variable(s.WithOpName("small"), {-1, 1, 1}, DT_FLOAT);
+  auto small = ops::Variable(s.WithOpName("small"), {-1, 1, 1}, DT_DOUBLE);
 
   // And second input which is larger, but has the same unknown dimension
   // device spec prevents this node from rewriting
-  auto d = "/job:do_not_rewrite_me";
-  auto v = ops::Variable(s.WithOpName("v"), {1, 32, 32}, DT_FLOAT);
+  auto d = "/device:CPU:0";
+  auto v = ops::Variable(s.WithOpName("v"), {1, 32, 32}, DT_DOUBLE);
   auto large = ops::Add(s.WithOpName("large").WithDevice(d), small, v);
 
   // [a, c] have {?, 1, 1} shape, [b] has {?, 32, 32}
@@ -1904,6 +1967,12 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MinimizeBCastWithSymbolicShapes) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto s_t = GenerateRandomTensor<DT_DOUBLE>(TensorShape({8, 1, 1}));
+  auto v_t = GenerateRandomTensor<DT_DOUBLE>(TensorShape({1, 32, 32}));
+  std::vector<std::pair<string, Tensor>> feed = {{"small", s_t}, {"v", v_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyAddToAddNCombining(&optimizer);
@@ -1942,6 +2011,10 @@ TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_MinimizeBCastWithSymbolicShapes) {
   const NodeDef* updated_outputs = node_map.GetNode("outputs");
   ASSERT_NE(updated_outputs, nullptr);
   EXPECT_EQ(outer_add_name, updated_outputs->input(0));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<double>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, RemoveNegation) {
@@ -1966,6 +2039,12 @@ TEST_F(ArithmeticOptimizerTest, RemoveNegation) {
   item.fetch = {"add_all"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto x_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  auto y_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 2}));
+  std::vector<std::pair<string, Tensor>> feed = {{"x", x_t}, {"y", y_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyRemoveNegation(&optimizer);
@@ -2014,6 +2093,10 @@ TEST_F(ArithmeticOptimizerTest, RemoveNegation) {
     }
   }
   EXPECT_EQ(5, found);
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, ConvertSqrtDivToRsqrtMul) {
@@ -2069,6 +2152,14 @@ TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_SimpleSwap) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto a_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32}));
+  auto b_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32, 32}));
+  auto c_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32}));
+  std::vector<std::pair<string, Tensor>> feed = {
+      {"a", a_t}, {"b", b_t}, {"c", c_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyMinimizeBroadcasts(&optimizer);
@@ -2093,16 +2184,20 @@ TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_SimpleSwap) {
   ASSERT_NE(mul2_node, nullptr);
   EXPECT_EQ("mul1", mul2_node->input(0));
   EXPECT_EQ("b", mul2_node->input(1));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_FlattenTallGraph) {
   tensorflow::Scope s = tensorflow::Scope::NewRootScope();
 
-  auto a = ops::Variable(s.WithOpName("a"), {32}, DT_FLOAT);
-  auto b = ops::Variable(s.WithOpName("b"), {32, 32}, DT_FLOAT);
-  auto c = ops::Variable(s.WithOpName("c"), {32}, DT_FLOAT);
-  auto d = ops::Variable(s.WithOpName("d"), {32}, DT_FLOAT);
-  auto e = ops::Variable(s.WithOpName("e"), {32}, DT_FLOAT);
+  auto a = ops::Variable(s.WithOpName("a"), {32}, DT_DOUBLE);
+  auto b = ops::Variable(s.WithOpName("b"), {32, 32}, DT_DOUBLE);
+  auto c = ops::Variable(s.WithOpName("c"), {32}, DT_DOUBLE);
+  auto d = ops::Variable(s.WithOpName("d"), {32}, DT_DOUBLE);
+  auto e = ops::Variable(s.WithOpName("e"), {32}, DT_DOUBLE);
 
   auto mul1 = ops::Mul(s.WithOpName("mul1"), a, b);
   auto mul2 = ops::Mul(s.WithOpName("mul2"), mul1, c);
@@ -2115,6 +2210,16 @@ TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_FlattenTallGraph) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto a_t = GenerateRandomTensor<DT_DOUBLE>(TensorShape({32}));
+  auto b_t = GenerateRandomTensor<DT_DOUBLE>(TensorShape({32, 32}));
+  auto c_t = GenerateRandomTensor<DT_DOUBLE>(TensorShape({32}));
+  auto d_t = GenerateRandomTensor<DT_DOUBLE>(TensorShape({32}));
+  auto e_t = GenerateRandomTensor<DT_DOUBLE>(TensorShape({32}));
+  std::vector<std::pair<string, Tensor>> feed = {
+      {"a", a_t}, {"b", b_t}, {"c", c_t}, {"d", d_t}, {"e", e_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyMinimizeBroadcasts(&optimizer);
@@ -2154,6 +2259,10 @@ TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_FlattenTallGraph) {
   ASSERT_NE(mul4_node, nullptr);
   EXPECT_EQ("mul3", mul4_node->input(0));
   EXPECT_EQ("b", mul4_node->input(1));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<double>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_BuildTreeUp) {
@@ -2175,6 +2284,15 @@ TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_BuildTreeUp) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto a_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32}));
+  auto b_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32}));
+  auto c_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32}));
+  auto d_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({32, 32}));
+  std::vector<std::pair<string, Tensor>> feed = {
+      {"a", a_t}, {"b", b_t}, {"c", c_t}, {"D", d_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyMinimizeBroadcasts(&optimizer);
@@ -2206,6 +2324,10 @@ TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_BuildTreeUp) {
   ASSERT_NE(mul3_node, nullptr);
   EXPECT_EQ("D", mul3_node->input(0));
   EXPECT_EQ("mul1", mul3_node->input(1));
+
+  auto tensors = EvaluateNodes(output, item.fetch, feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryFromConcat) {