| 409 | } |
| 410 | |
| 411 | StatusOr<bool> ScatterExpander::Run(HloModule* module) { |
| 412 | std::vector<HloInstruction*> scatter_instrs; |
| 413 | for (HloComputation* computation : module->MakeNonfusionComputations()) { |
| 414 | for (HloInstruction* instr : computation->instructions()) { |
| 415 | if (instr->opcode() == HloOpcode::kScatter) { |
| 416 | scatter_instrs.push_back(instr); |
| 417 | } |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | for (auto instr : scatter_instrs) { |
| 422 | TF_ASSIGN_OR_RETURN(HloInstruction * expanded_root, ExpandScatter(instr)); |
| 423 | TF_RETURN_IF_ERROR( |
| 424 | instr->parent()->ReplaceInstruction(instr, expanded_root)); |
| 425 | } |
| 426 | |
| 427 | return !scatter_instrs.empty(); |
| 428 | } |
| 429 | |
| 430 | } // namespace xla |
nothing calls this directly
no test coverage detected