Find the scf.if that guards the epilogue (masking + sample + store). This is the scf.if whose body contains a tt.reduce (the argmax).
(lines: list[str])
| 147 | |
| 148 | |
| 149 | def _find_epilogue_if(lines: list[str]) -> int | None: |
| 150 | """Find the scf.if that guards the epilogue (masking + sample + store). |
| 151 | |
| 152 | This is the scf.if whose body contains a tt.reduce (the argmax). |
| 153 | """ |
| 154 | for i, line in enumerate(lines): |
| 155 | if "scf.if" not in line: |
| 156 | continue |
| 157 | # Look ahead to see if this scf.if contains a tt.reduce |
| 158 | for j in range(i + 1, min(i + 300, len(lines))): |
| 159 | if '"tt.reduce"' in lines[j]: |
| 160 | return i |
| 161 | # Stop if we hit the matching else/closing brace at the same nesting |
| 162 | if "} else {" in lines[j] or lines[j].strip().startswith("} loc("): |
| 163 | break |
| 164 | return None |
| 165 | |
| 166 | |
| 167 | def _find_reduce_close(lines: list[str]) -> int | None: |
no outgoing calls
no test coverage detected