A `Dataset` that maps a function over its input and interleaves the result.
| 3661 | |
| 3662 | |
| 3663 | class ParallelInterleaveDataset(UnaryDataset): |
| 3664 | """A `Dataset` that maps a function over its input and interleaves the result.""" |
| 3665 | |
| 3666 | def __init__(self, input_dataset, map_func, cycle_length, block_length, |
| 3667 | num_parallel_calls): |
| 3668 | """See `Dataset.interleave()` for details.""" |
| 3669 | self._input_dataset = input_dataset |
| 3670 | self._map_func = StructuredFunctionWrapper( |
| 3671 | map_func, self._transformation_name(), dataset=input_dataset) |
| 3672 | if not isinstance(self._map_func.output_structure, DatasetSpec): |
| 3673 | raise TypeError( |
| 3674 | "`map_func` must return a `Dataset` object. Got {}".format( |
| 3675 | type(self._map_func.output_structure))) |
| 3676 | self._structure = self._map_func.output_structure._element_spec # pylint: disable=protected-access |
| 3677 | self._cycle_length = ops.convert_to_tensor( |
| 3678 | cycle_length, dtype=dtypes.int64, name="cycle_length") |
| 3679 | self._block_length = ops.convert_to_tensor( |
| 3680 | block_length, dtype=dtypes.int64, name="block_length") |
| 3681 | self._num_parallel_calls = ops.convert_to_tensor( |
| 3682 | num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls") |
| 3683 | variant_tensor = gen_dataset_ops.parallel_interleave_dataset_v2( |
| 3684 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3685 | self._map_func.function.captured_inputs, # pylint: disable=protected-access |
| 3686 | self._cycle_length, |
| 3687 | self._block_length, |
| 3688 | self._num_parallel_calls, |
| 3689 | f=self._map_func.function, |
| 3690 | **self._flat_structure) |
| 3691 | super(ParallelInterleaveDataset, self).__init__(input_dataset, |
| 3692 | variant_tensor) |
| 3693 | |
| 3694 | def _functions(self): |
| 3695 | return [self._map_func] |
| 3696 | |
| 3697 | @property |
| 3698 | def element_spec(self): |
| 3699 | return self._structure |
| 3700 | |
| 3701 | def _transformation_name(self): |
| 3702 | return "Dataset.interleave()" |
| 3703 | |
| 3704 | |
| 3705 | class FilterDataset(UnaryUnchangedStructureDataset): |