Function TEST

tensorflow/compiler/tf2xla/const_analysis_test.cc:31–60

namespace {

TEST(ConstAnalysisTest, Basics) {
  Scope root = Scope::NewRootScope();

  auto arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0);
  auto arg1 = ops::_Arg(root.WithOpName("Arg1"), DT_INT32, 1);
  auto arg2 = ops::_Arg(root.WithOpName("Arg2"), DT_INT32, 2);
  auto arg3 = ops::_Arg(root.WithOpName("Arg3"), DT_INT32, 3);
  auto a = ops::Shape(root, arg0);
  auto b = ops::Add(root, a, arg1);
  auto c = ops::Reshape(root, arg2, b);
  auto d = ops::Mul(root, c, ops::Sum(root, arg3, arg3));

  FixupSourceAndSinkEdges(root.graph());

  std::vector<bool> const_args(4, false);
  std::vector<bool> const_nodes(root.graph()->num_node_ids(), false);
  TF_ASSERT_OK(BackwardsConstAnalysis(*root.graph(), &const_args, &const_nodes,
                                      /*flib_runtime=*/nullptr));

  // Arg 0 doesn't need to be constant since the graph only uses its shape.
  // Arg 1 must be constant because it flows to the shape argument of a Reshape.
  // Arg 2 is used only as the value input to a Reshape and need not be const.
  // Arg 3 is used as the reduction-indices argument to Sum and must be const.
  EXPECT_EQ(const_args, std::vector<bool>({false, true, false, true}));

  EXPECT_FALSE(const_nodes[arg0.node()->id()]);
  EXPECT_TRUE(const_nodes[arg1.node()->id()]);
  EXPECT_FALSE(const_nodes[arg2.node()->id()]);
  EXPECT_TRUE(const_nodes[arg3.node()->id()]);
}

// Regression test for a case where the backward const analysis did
// not visit nodes in topological order.