A `Dataset` that maps a function over its input and flattens the result.
| 220 | |
| 221 | |
| 222 | class ParallelInterleaveDataset(dataset_ops.UnaryDataset): |
| 223 | """A `Dataset` that maps a function over its input and flattens the result.""" |
| 224 | |
| 225 | def __init__(self, input_dataset, map_func, cycle_length, block_length, |
| 226 | sloppy, buffer_output_elements, prefetch_input_elements): |
| 227 | """See `tf.data.experimental.parallel_interleave()` for details.""" |
| 228 | self._input_dataset = input_dataset |
| 229 | self._map_func = dataset_ops.StructuredFunctionWrapper( |
| 230 | map_func, self._transformation_name(), dataset=input_dataset) |
| 231 | if not isinstance(self._map_func.output_structure, dataset_ops.DatasetSpec): |
| 232 | raise TypeError("`map_func` must return a `Dataset` object.") |
| 233 | self._element_spec = self._map_func.output_structure._element_spec # pylint: disable=protected-access |
| 234 | self._cycle_length = ops.convert_to_tensor( |
| 235 | cycle_length, dtype=dtypes.int64, name="cycle_length") |
| 236 | self._block_length = ops.convert_to_tensor( |
| 237 | block_length, dtype=dtypes.int64, name="block_length") |
| 238 | self._sloppy = ops.convert_to_tensor( |
| 239 | sloppy, dtype=dtypes.bool, name="sloppy") |
| 240 | self._buffer_output_elements = convert.optional_param_to_tensor( |
| 241 | "buffer_output_elements", |
| 242 | buffer_output_elements, |
| 243 | argument_default=2 * block_length) |
| 244 | self._prefetch_input_elements = convert.optional_param_to_tensor( |
| 245 | "prefetch_input_elements", |
| 246 | prefetch_input_elements, |
| 247 | argument_default=2 * cycle_length) |
| 248 | variant_tensor = ged_ops.parallel_interleave_dataset( |
| 249 | self._input_dataset._variant_tensor, # pylint: disable=protected-access |
| 250 | self._map_func.function.captured_inputs, |
| 251 | self._cycle_length, |
| 252 | self._block_length, |
| 253 | self._sloppy, |
| 254 | self._buffer_output_elements, |
| 255 | self._prefetch_input_elements, |
| 256 | f=self._map_func.function, |
| 257 | **self._flat_structure) |
| 258 | super(ParallelInterleaveDataset, self).__init__(input_dataset, |
| 259 | variant_tensor) |
| 260 | |
| 261 | def _functions(self): |
| 262 | return [self._map_func] |
| 263 | |
| 264 | @property |
| 265 | def element_spec(self): |
| 266 | return self._element_spec |
| 267 | |
| 268 | def _transformation_name(self): |
| 269 | return "tf.data.experimental.parallel_interleave()" |
| 270 | |
| 271 | |
| 272 | @tf_export("data.TFRecordDataset", v1=[]) |
no outgoing calls
no test coverage detected