(*inner_args, **inner_kwargs)
| 501 | ix = [0] |
| 502 | |
| 503 | def enumerated_fn(*inner_args, **inner_kwargs): |
| 504 | r = map_fn(ix[0], *inner_args, **inner_kwargs) |
| 505 | ix[0] += 1 |
| 506 | return r |
| 507 | |
| 508 | return nest.map_structure_up_to(shallow_structure, enumerated_fn, *args, |
| 509 | **kwargs) |