A `Dataset` that acts as an identity, and models performance.
| 3820 | |
| 3821 | |
| 3822 | class _ModelDataset(UnaryUnchangedStructureDataset): |
| 3823 | """A `Dataset` that acts as an identity, and models performance.""" |
| 3824 | |
| 3825 | def __init__(self, input_dataset, algorithm, cpu_budget): |
| 3826 | self._input_dataset = input_dataset |
| 3827 | # TODO(jsimsa): This check is introduced for forward compatibility and can |
| 3828 | # be removed after 7/24/2019. At that point, all servers are expected to |
| 3829 | # recognize the `algorithm` attribute. |
| 3830 | if algorithm != AutotuneAlgorithm.HILL_CLIMB: |
| 3831 | variant_tensor = gen_dataset_ops.model_dataset( |
| 3832 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3833 | algorithm=algorithm, |
| 3834 | cpu_budget=cpu_budget, |
| 3835 | **self._flat_structure) |
| 3836 | else: |
| 3837 | variant_tensor = gen_dataset_ops.model_dataset( |
| 3838 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3839 | cpu_budget=cpu_budget, |
| 3840 | **self._flat_structure) |
| 3841 | super(_ModelDataset, self).__init__(input_dataset, variant_tensor) |
| 3842 | |
| 3843 | |
| 3844 | class _OptimizeDataset(UnaryUnchangedStructureDataset): |