Function TEST_F

tensorflow/compiler/jit/xla_kernel_creator_test.cc:95–120

}

TEST_F(XlaKernelCreatorTest, OneFloatOneResourceArgument) {
  FunctionDef fdef = XTimesY();
  (*fdef.mutable_attr())["_XlaCompile"] = BoolAttr(true);
  Init({fdef});
  XlaKernelCreator xla_kernel_creator;

  Status status = xla_kernel_creator.CreateKernel(
      flr_, ToNodeDef(R"pb(
        name: 'XTimesY' op: 'XTimesY' input: 'a' input: 'b'
      )pb"),
      &kernel_);
  ASSERT_TRUE(status.ok()) << status.ToString();

  EXPECT_EQ("XTimesY", kernel_->name());
  EXPECT_EQ("XTimesY", kernel_->type_string());

  EXPECT_EQ(2, kernel_->num_inputs());
  EXPECT_EQ(DT_FLOAT, kernel_->input_type(0));
  EXPECT_EQ(DT_RESOURCE, kernel_->input_type(1));
  EXPECT_EQ(DEVICE_MEMORY, kernel_->input_memory_types()[0]);
  EXPECT_EQ(HOST_MEMORY, kernel_->input_memory_types()[1]);

  EXPECT_EQ(1, kernel_->num_outputs());
  EXPECT_EQ(DT_FLOAT, kernel_->output_type(0));
  EXPECT_EQ(DEVICE_MEMORY, kernel_->output_memory_types()[0]);
}

TEST_F(XlaKernelCreatorTest, FailsIfXlaCompileAttrNotSet) {
  FunctionDef fdef = XTimesY();