| 108 | } |
| 109 | |
| 110 | StatusOr<bool> BatchDotSimplification::Run(HloModule* module) { |
| 111 | bool changed = false; |
| 112 | std::vector<HloInstruction*> dot_instrs; |
| 113 | for (HloComputation* computation : module->MakeNonfusionComputations()) { |
| 114 | absl::c_copy_if(computation->instructions(), std::back_inserter(dot_instrs), |
| 115 | [](HloInstruction* instr) { |
| 116 | return instr->opcode() == HloOpcode::kDot; |
| 117 | }); |
| 118 | } |
| 119 | for (HloInstruction* dot_instr : dot_instrs) { |
| 120 | TF_ASSIGN_OR_RETURN(bool elided_batch_dim_from_one, |
| 121 | ElideDegenerateBatchDimensionFromBatchDot(dot_instr)); |
| 122 | changed |= elided_batch_dim_from_one; |
| 123 | } |
| 124 | return changed; |
| 125 | } |
| 126 | } // namespace xla |
nothing calls this directly
no test coverage detected