An internal helper for changing the structure and shape of a dataset.
| 3911 | |
| 3912 | |
| 3913 | class _RestructuredDataset(UnaryDataset): |
| 3914 | """An internal helper for changing the structure and shape of a dataset.""" |
| 3915 | |
| 3916 | def __init__(self, dataset, structure): |
| 3917 | self._input_dataset = dataset |
| 3918 | self._structure = structure |
| 3919 | |
| 3920 | variant_tensor = self._input_dataset._variant_tensor # pylint: disable=protected-access |
| 3921 | super(_RestructuredDataset, self).__init__(dataset, variant_tensor) |
| 3922 | |
| 3923 | @property |
| 3924 | def element_spec(self): |
| 3925 | return self._structure |
| 3926 | |
| 3927 | |
| 3928 | class _UnbatchDataset(UnaryDataset): |