| 34 | }; |
| 35 | |
| 36 | TEST_F(HloModuleGroupTest, SingleModule) { |
| 37 | const string text = R"( |
| 38 | HloModule simple_module |
| 39 | |
| 40 | ENTRY %entry (x: f32[], y: f32[]) -> f32[] { |
| 41 | %x = f32[] parameter(0) |
| 42 | %y = f32[] parameter(1) |
| 43 | ROOT %add = f32[] add(%x, %y) |
| 44 | } |
| 45 | )"; |
| 46 | TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module, |
| 47 | ParseAndReturnVerifiedModule(text)); |
| 48 | HloModuleGroup group(std::move(module)); |
| 49 | |
| 50 | EXPECT_EQ(group.modules().size(), 1); |
| 51 | EXPECT_THAT( |
| 52 | group.module(0).entry_computation()->instructions(), |
| 53 | ::testing::ElementsAre(op::Parameter(), op::Parameter(), op::Add())); |
| 54 | |
| 55 | TF_ASSERT_OK_AND_ASSIGN(HloModuleGroup group_copy, |
| 56 | HloModuleGroup::CreateFromProto( |
| 57 | group.ToProto(), {group.module(0).config()})); |
| 58 | EXPECT_EQ(group_copy.modules().size(), 1); |
| 59 | EXPECT_THAT( |
| 60 | group_copy.module(0).entry_computation()->instructions(), |
| 61 | ::testing::ElementsAre(op::Parameter(), op::Parameter(), op::Add())); |
| 62 | |
| 63 | std::vector<std::unique_ptr<HloModule>> modules = group.ConsumeModules(); |
| 64 | EXPECT_EQ(modules.size(), 1); |
| 65 | EXPECT_EQ(group.modules().size(), 0); |
| 66 | } |
| 67 | |
| 68 | TEST_F(HloModuleGroupTest, MultipleModules) { |
| 69 | const string text_0 = R"( |
nothing calls this directly
no test coverage detected