(i)
| 384 | ``` |
| 385 | """ |
| 386 | def loop_fn(i): |
| 387 | gathered_elems = nest.map_structure(lambda x: array_ops.gather(x, i), elems) |
| 388 | return fn(gathered_elems) |
| 389 | batch_size = array_ops.shape(nest.flatten(elems)[0])[0] |
| 390 | return pfor(loop_fn, batch_size) |
no test coverage detected