Function TEST_F

tensorflow/compiler/xla/service/map_inliner_test.cc:41–75

Test that `map` with `max` is transformed to `max`

// Test that `map` with `max` is transformed to `max`
TEST_F(MapInlinerTest, MapMax) {
  Shape r0f32 = ShapeUtil::MakeShape(F32, {});

  auto max_builder = HloComputation::Builder(TestName());
  auto param1 = max_builder.AddInstruction(
      HloInstruction::CreateParameter(0, r0f32, "x"));
  auto param2 = max_builder.AddInstruction(
      HloInstruction::CreateParameter(1, r0f32, "y"));
  max_builder.AddInstruction(HloInstruction::CreateBinary(
      param1->shape(), HloOpcode::kMaximum, param1, param2));
  auto max_f32 = max_builder.Build();

  auto builder = HloComputation::Builder("MapMaxFunction");
  auto lhs = builder.AddInstruction(HloInstruction::CreateConstant(
      LiteralUtil::CreateR1<float>({1, 2, 3, 4})));
  auto rhs = builder.AddInstruction(HloInstruction::CreateConstant(
      LiteralUtil::CreateR1<float>({4, 3, 2, 1})));
  builder.AddInstruction(
      HloInstruction::CreateMap(lhs->shape(), {lhs, rhs}, max_f32.get()));

  auto computation = builder.Build();
  auto hlo_module = CreateNewVerifiedModule();
  hlo_module->AddEmbeddedComputation(std::move(max_f32));
  hlo_module->AddEntryComputation(std::move(computation));

  MapInliner inliner;
  EXPECT_TRUE(inliner.Run(hlo_module.get()).ValueOrDie());
  EXPECT_THAT(hlo_module->entry_computation()->root_instruction(),
              op::Maximum(lhs, rhs));

  // Verify execution on CPU.
  auto result = ExecuteAndTransfer(hlo_module->Clone(), {});
  auto expected = LiteralUtil::CreateR1<float>({4, 3, 3, 4});
  EXPECT_TRUE(LiteralTestUtil::Equal(result, expected));
}

// Test that `constant` function is changed to `broadcast`.
TEST_F(MapInlinerTest, MapConstant) {