Executes `reduce_fn` in a loop till the dataset is empty.
(has_data, data, state)
| 426 | return has_data |
| 427 | |
| 428 | def loop_body(has_data, data, state): |
| 429 | """Executes `reduce_fn` in a loop till the dataset is empty.""" |
| 430 | del has_data # Unused. |
| 431 | # data is list of lists here. where each list corresponds to one worker. |
| 432 | # TODO(b/130570614): Add support for the multiworker and TPU pods use |
| 433 | # case. |
| 434 | if self._input_workers.num_workers == 1: |
| 435 | data = data[0] |
| 436 | else: |
| 437 | raise ValueError("Dataset iteration within a tf.function is" |
| 438 | " not supported for multiple workers.") |
| 439 | per_replica_data = values.regroup(self._input_workers.device_map, data) |
| 440 | state = reduce_fn(state, per_replica_data) |
| 441 | has_data, data = _get_next_as_optional(iterator, self._strategy) |
| 442 | return has_data, data, state |
| 443 | |
| 444 | has_data, data, final_state = control_flow_ops.while_loop( |
| 445 | cond, loop_body, [has_data, data, initial_state], parallel_iterations=1) |
nothing calls this directly
no test coverage detected