| 669 | } |
| 670 | |
| 671 | Status HloModule::RemoveUnusedComputations() { |
| 672 | std::string suffix = "tmp"; |
| 673 | auto module = absl::make_unique<HloModule>( |
| 674 | absl::StrCat(name_, suffix.empty() ? "" : "-", suffix), config()); |
| 675 | HloCloneContext context(module.get(), suffix); |
| 676 | entry_computation_->Clone(suffix, &context); |
| 677 | std::vector<HloComputation*> to_remove; |
| 678 | for (auto computation : computations()) { |
| 679 | auto found_computation = context.FindComputation(computation); |
| 680 | if (found_computation == nullptr) { |
| 681 | to_remove.push_back(computation); |
| 682 | } |
| 683 | } |
| 684 | for (auto computation : to_remove) { |
| 685 | TF_RETURN_IF_ERROR(RemoveEmbeddedComputation(computation)); |
| 686 | } |
| 687 | return Status::OK(); |
| 688 | } |
| 689 | |
| 690 | HloComputation* HloModule::DeepCloneComputation(HloComputation* computation, |
| 691 | HloCloneContext* context) { |