(i, pfor_config=None)
| 280 | offset = j * parallel_iterations + num_remaining_iterations |
| 281 | |
| 282 | def tiled_loop_fn(i, pfor_config=None): |
| 283 | if loop_fn_has_config: |
| 284 | return nest.flatten(loop_fn(i + offset, pfor_config=pfor_config)) |
| 285 | else: |
| 286 | return nest.flatten(loop_fn(i + offset)) |
| 287 | |
| 288 | return _pfor_impl( |
| 289 | tiled_loop_fn, parallel_iterations, pfor_config=pfor_config) |