Function TEST

tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc:86–143

}

TEST(XlaJitCompiledCpuFunction, Sum) {
  GraphDef graph_def = SumGraph();
  tf2xla::Config config = SumConfig();

  TF_ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<XlaJitCompiledCpuFunction> jit,
      XlaJitCompiledCpuFunction::Compile(graph_def, config,
                                         xla::ExecutableBuildOptions()));
  XlaCompiledCpuFunction function(jit->StaticData());

  // Run the function and check results.
  *static_cast<int32*>(function.arg_data(0)) = 10;
  *static_cast<int32*>(function.arg_data(1)) = 32;
  EXPECT_TRUE(function.Run());
  EXPECT_EQ(function.error_msg(), "");
  EXPECT_EQ(*static_cast<int32*>(function.result_data(0)), 42);

  // Run the function again.
  *static_cast<int32*>(function.arg_data(0)) = 100;
  *static_cast<int32*>(function.arg_data(1)) = 320;
  EXPECT_TRUE(function.Run());
  EXPECT_EQ(function.error_msg(), "");
  EXPECT_EQ(*static_cast<int32*>(function.result_data(0)), 420);

  // Check name to index lookups.
  EXPECT_TRUE(function.HasNameIndices());

  EXPECT_EQ(function.LookupArgIndex("x_name"), 0);
  EXPECT_EQ(function.LookupArgIndex("y_name"), 1);
  EXPECT_EQ(function.LookupArgIndex(""), -1);
  EXPECT_EQ(function.LookupArgIndex("x"), -1);
  EXPECT_EQ(function.LookupArgIndex("y"), -1);
  EXPECT_EQ(function.LookupArgIndex("sum"), -1);
  EXPECT_EQ(function.LookupArgIndex("sum_name"), -1);

  EXPECT_EQ(function.LookupResultIndex("sum_name"), 0);
  EXPECT_EQ(function.LookupResultIndex(""), -1);
  EXPECT_EQ(function.LookupResultIndex("x"), -1);
  EXPECT_EQ(function.LookupResultIndex("y"), -1);
  EXPECT_EQ(function.LookupResultIndex("sum"), -1);
  EXPECT_EQ(function.LookupResultIndex("x_name"), -1);
  EXPECT_EQ(function.LookupResultIndex("y_name"), -1);

  // Check program shape.
  using xla::ShapeUtil;
  const xla::Shape s32 = ShapeUtil::MakeShape(xla::S32, {});
  ASSERT_TRUE(function.ProgramShape() != nullptr);
  const xla::ProgramShape program_shape(*function.ProgramShape());
  ASSERT_EQ(program_shape.parameters_size(), 2);
  EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(0), s32));
  EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(1), s32));

  const xla::Shape& result = program_shape.result();
  ASSERT_EQ(result.element_type(), xla::TUPLE);
  ASSERT_EQ(ShapeUtil::TupleElementCount(result), 1);
  const xla::Shape& result0 = ShapeUtil::GetTupleElementShape(result, 0);
  EXPECT_TRUE(ShapeUtil::Compatible(result0, s32));
}