A `Dataset` that interleaves the result of transformed inputs.
| 3623 | |
| 3624 | |
| 3625 | class InterleaveDataset(UnaryDataset): |
| 3626 | """A `Dataset` that interleaves the result of transformed inputs.""" |
| 3627 | |
| 3628 | def __init__(self, input_dataset, map_func, cycle_length, block_length): |
| 3629 | """See `Dataset.interleave()` for details.""" |
| 3630 | self._input_dataset = input_dataset |
| 3631 | self._map_func = StructuredFunctionWrapper( |
| 3632 | map_func, self._transformation_name(), dataset=input_dataset) |
| 3633 | if not isinstance(self._map_func.output_structure, DatasetSpec): |
| 3634 | raise TypeError( |
| 3635 | "`map_func` must return a `Dataset` object. Got {}".format( |
| 3636 | type(self._map_func.output_structure))) |
| 3637 | self._structure = self._map_func.output_structure._element_spec # pylint: disable=protected-access |
| 3638 | self._cycle_length = ops.convert_to_tensor( |
| 3639 | cycle_length, dtype=dtypes.int64, name="cycle_length") |
| 3640 | self._block_length = ops.convert_to_tensor( |
| 3641 | block_length, dtype=dtypes.int64, name="block_length") |
| 3642 | |
| 3643 | variant_tensor = gen_dataset_ops.interleave_dataset( |
| 3644 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3645 | self._map_func.function.captured_inputs, # pylint: disable=protected-access |
| 3646 | self._cycle_length, |
| 3647 | self._block_length, |
| 3648 | f=self._map_func.function, |
| 3649 | **self._flat_structure) |
| 3650 | super(InterleaveDataset, self).__init__(input_dataset, variant_tensor) |
| 3651 | |
| 3652 | def _functions(self): |
| 3653 | return [self._map_func] |
| 3654 | |
| 3655 | @property |
| 3656 | def element_spec(self): |
| 3657 | return self._structure |
| 3658 | |
| 3659 | def _transformation_name(self): |
| 3660 | return "Dataset.interleave()" |
| 3661 | |
| 3662 | |
| 3663 | class ParallelInterleaveDataset(UnaryDataset): |