Type specification for `tf.data.Dataset`.
| 2576 | "data.DatasetSpec", |
| 2577 | v1=["data.DatasetSpec", "data.experimental.DatasetStructure"]) |
| 2578 | class DatasetSpec(type_spec.BatchableTypeSpec): |
| 2579 | """Type specification for `tf.data.Dataset`.""" |
| 2580 | |
| 2581 | __slots__ = ["_element_spec", "_dataset_shape"] |
| 2582 | |
| 2583 | def __init__(self, element_spec, dataset_shape=()): |
| 2584 | self._element_spec = element_spec |
| 2585 | self._dataset_shape = tensor_shape.as_shape(dataset_shape) |
| 2586 | |
| 2587 | @property |
| 2588 | def value_type(self): |
| 2589 | return _VariantDataset |
| 2590 | |
| 2591 | def _serialize(self): |
| 2592 | return (self._element_spec, self._dataset_shape) |
| 2593 | |
| 2594 | @property |
| 2595 | def _component_specs(self): |
| 2596 | return tensor_spec.TensorSpec(self._dataset_shape, dtypes.variant) |
| 2597 | |
| 2598 | def _to_components(self, value): |
| 2599 | return value._variant_tensor # pylint: disable=protected-access |
| 2600 | |
| 2601 | def _from_components(self, components): |
| 2602 | # pylint: disable=protected-access |
| 2603 | if self._dataset_shape.ndims == 0: |
| 2604 | return _VariantDataset(components, self._element_spec) |
| 2605 | else: |
| 2606 | return _NestedVariant(components, self._element_spec, self._dataset_shape) |
| 2607 | |
| 2608 | def _to_tensor_list(self, value): |
| 2609 | return [ |
| 2610 | ops.convert_to_tensor( |
| 2611 | tf_nest.map_structure(lambda x: x._variant_tensor, value)) # pylint: disable=protected-access |
| 2612 | ] |
| 2613 | |
| 2614 | @staticmethod |
| 2615 | def from_value(value): |
| 2616 | return DatasetSpec(value.element_spec) # pylint: disable=protected-access |
| 2617 | |
| 2618 | def _batch(self, batch_size): |
| 2619 | return DatasetSpec( |
| 2620 | self._element_spec, |
| 2621 | tensor_shape.TensorShape([batch_size]).concatenate(self._dataset_shape)) |
| 2622 | |
| 2623 | def _unbatch(self): |
| 2624 | if self._dataset_shape.ndims == 0: |
| 2625 | raise ValueError("Unbatching a dataset is only supported for rank >= 1") |
| 2626 | return DatasetSpec(self._element_spec, self._dataset_shape[1:]) |
| 2627 | |
| 2628 | def _to_batched_tensor_list(self, value): |
| 2629 | if self._dataset_shape.ndims == 0: |
| 2630 | raise ValueError("Unbatching a dataset is only supported for rank >= 1") |
| 2631 | return self._to_tensor_list(value) |
| 2632 | |
| 2633 | def _to_legacy_output_types(self): |
| 2634 | return self |
| 2635 |
no outgoing calls
no test coverage detected