A `Dataset` that maps a function over its input and flattens the result.
| 3592 | |
| 3593 | |
| 3594 | class FlatMapDataset(UnaryDataset): |
| 3595 | """A `Dataset` that maps a function over its input and flattens the result.""" |
| 3596 | |
| 3597 | def __init__(self, input_dataset, map_func): |
| 3598 | """See `Dataset.flat_map()` for details.""" |
| 3599 | self._input_dataset = input_dataset |
| 3600 | self._map_func = StructuredFunctionWrapper( |
| 3601 | map_func, self._transformation_name(), dataset=input_dataset) |
| 3602 | if not isinstance(self._map_func.output_structure, DatasetSpec): |
| 3603 | raise TypeError( |
| 3604 | "`map_func` must return a `Dataset` object. Got {}".format( |
| 3605 | type(self._map_func.output_structure))) |
| 3606 | self._structure = self._map_func.output_structure._element_spec # pylint: disable=protected-access |
| 3607 | variant_tensor = gen_dataset_ops.flat_map_dataset( |
| 3608 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3609 | self._map_func.function.captured_inputs, |
| 3610 | f=self._map_func.function, |
| 3611 | **self._flat_structure) |
| 3612 | super(FlatMapDataset, self).__init__(input_dataset, variant_tensor) |
| 3613 | |
| 3614 | def _functions(self): |
| 3615 | return [self._map_func] |
| 3616 | |
| 3617 | @property |
| 3618 | def element_spec(self): |
| 3619 | return self._structure |
| 3620 | |
| 3621 | def _transformation_name(self): |
| 3622 | return "Dataset.flat_map()" |
| 3623 | |
| 3624 | |
| 3625 | class InterleaveDataset(UnaryDataset): |