A `Dataset` that acts as an identity, and sets a stats aggregator.
| 3861 | |
| 3862 | |
| 3863 | class _SetStatsAggregatorDataset(UnaryUnchangedStructureDataset): |
| 3864 | """A `Dataset` that acts as an identity, and sets a stats aggregator.""" |
| 3865 | |
| 3866 | def __init__(self, input_dataset, aggregator, prefix, counter_prefix): |
| 3867 | self._input_dataset = input_dataset |
| 3868 | self._stats_aggregator = aggregator |
| 3869 | self._prefix = prefix |
| 3870 | self._counter_prefix = counter_prefix |
| 3871 | variant_tensor = ged_ops.set_stats_aggregator_dataset( |
| 3872 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3873 | self._stats_aggregator._resource, # pylint: disable=protected-access |
| 3874 | self._prefix, |
| 3875 | self._counter_prefix, |
| 3876 | **self._flat_structure) |
| 3877 | super(_SetStatsAggregatorDataset, self).__init__(input_dataset, |
| 3878 | variant_tensor) |
| 3879 | |
| 3880 | |
| 3881 | class _MaxIntraOpParallelismDataset(UnaryUnchangedStructureDataset): |