| 135 | } // namespace |
| 136 | |
| 137 | std::unique_ptr<HloModule> ExtractModule(HloInstruction* instruction, |
| 138 | int64 height) { |
| 139 | absl::flat_hash_set<const HloInstruction*> boundary; |
| 140 | if (height != -1) { |
| 141 | ComputeBoundary(instruction, height, &boundary); |
| 142 | } |
| 143 | ExtractionVisitor visitor(*instruction->GetModule(), &boundary); |
| 144 | CHECK(instruction->Accept(&visitor).ok()); |
| 145 | |
| 146 | // The first pass may leave unused parameter instructions. Do another |
| 147 | // extraction pass to remove unused parameters. This is done because |
| 148 | // HloComputation does not allow removing parameters after the computation has |
| 149 | // been built. |
| 150 | ExtractionVisitor cleanup_visitor(*visitor.module(), /*boundary=*/nullptr); |
| 151 | TF_CHECK_OK(visitor.module()->entry_computation()->root_instruction()->Accept( |
| 152 | &cleanup_visitor)); |
| 153 | |
| 154 | HloVerifier verifier(/*layout_sensitive=*/false, |
| 155 | /*allow_mixed_precision=*/true); |
| 156 | TF_CHECK_OK(verifier.Run(cleanup_visitor.module()).status()); |
| 157 | return cleanup_visitor.ConsumeModule(); |
| 158 | } |
| 159 | |
| 160 | } // namespace xla |