Function TEST_F

tensorflow/compiler/xla/service/conditional_to_select_test.cc:39–78

Test that a conditional of simple constants is transformed to a select

// Test that a conditional of simple constants is transformed to a select
TEST_F(ConditionalToSelectTest, MapConditionalConstants) {
  const string hlo_text = R"(
HloModule MapConditionalConstants

if {
  %pif = () parameter(0)
  ROOT %cif = f32[] constant(0)
}

else {
  %pelse = () parameter(0)
  ROOT %celse = f32[] constant(1)
}

mapped {
  %a = f32[] parameter(0)
  %b = f32[] parameter(1)
  %lt = pred[] compare(%a, %b), direction=LT
  %t = () tuple()
  ROOT %conditional = f32[] conditional(%lt, %t, %t), true_computation=if, false_computation=else
}

ENTRY comp {
  %p1 = f32[1000]{0} parameter(0)
  %p2 = f32[1000]{0} parameter(1)
  ROOT %mapped = f32[1000]{0} map(%p1, %p2), dimensions={0}, to_apply=mapped
}
)";

  auto module = ParseAndReturnVerifiedModule(hlo_text).ValueOrDie();
  ConditionalToSelect pass;
  ASSERT_TRUE(pass.Run(&*module).ValueOrDie());

  HloInstruction* root = module->entry_computation()->root_instruction();
  ASSERT_EQ(root->opcode(), HloOpcode::kMap);
  HloComputation* mapped = root->called_computations()[0];
  EXPECT_THAT(mapped->root_instruction(),
              op::Select(op::Lt(op::Parameter(0), op::Parameter(1)),
                         op::Constant(), op::Constant()));
}

// Test that the condition gets broadcasted for feeding into
// select when the output is non-scalar.