| 26 | using HloExtractorTest = HloTestBase; |
| 27 | |
| 28 | TEST_F(HloExtractorTest, ExtractTopLevel) { |
| 29 | const string& hlo_string = R"( |
| 30 | HloModule test |
| 31 | |
| 32 | ENTRY %entry { |
| 33 | param.0 = f32[4]{0} parameter(0) |
| 34 | negate = f32[4]{0} negate(f32[4]{0} param.0) |
| 35 | ROOT exp = f32[4]{0} exponential(f32[4]{0} negate) |
| 36 | } |
| 37 | )"; |
| 38 | |
| 39 | TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module, |
| 40 | ParseAndReturnVerifiedModule(hlo_string)); |
| 41 | |
| 42 | { |
| 43 | auto extracted_module = |
| 44 | ExtractModule(FindInstruction(hlo_module.get(), "exp")); |
| 45 | EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), |
| 46 | op::Exp(op::Negate(op::Parameter(0)))); |
| 47 | } |
| 48 | |
| 49 | { |
| 50 | auto extracted_module = |
| 51 | ExtractModule(FindInstruction(hlo_module.get(), "exp"), /*height=*/0); |
| 52 | EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), |
| 53 | op::Exp(op::Parameter(0))); |
| 54 | } |
| 55 | |
| 56 | { |
| 57 | auto extracted_module = ExtractModule( |
| 58 | FindInstruction(hlo_module.get(), "negate"), /*height=*/0); |
| 59 | EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), |
| 60 | op::Negate(op::Parameter(0))); |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | TEST_F(HloExtractorTest, ExtractDag) { |
| 65 | const string& hlo_string = R"( |
nothing calls this directly
no test coverage detected