Execute a `reduce_fn` over all the elements of the input.
(self, initial_state, reduce_fn)
| 417 | return () |
| 418 | |
| 419 | def reduce(self, initial_state, reduce_fn): |
| 420 | """Execute a `reduce_fn` over all the elements of the input.""" |
| 421 | iterator = iter(self) |
| 422 | has_data, data = _get_next_as_optional(iterator, self._strategy) |
| 423 | |
| 424 | def cond(has_data, data, state): |
| 425 | del data, state # Unused. |
| 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) |
| 446 | return final_state |
| 447 | |
| 448 | |
| 449 | class DistributedDataset(_IterableInput): |
no test coverage detected