A `Dataset` that acts as an identity, setting a private threadpool.
| 3896 | |
| 3897 | |
| 3898 | class _PrivateThreadPoolDataset(UnaryUnchangedStructureDataset): |
| 3899 | """A `Dataset` that acts as an identity, setting a private threadpool.""" |
| 3900 | |
| 3901 | def __init__(self, input_dataset, num_threads): |
| 3902 | self._input_dataset = input_dataset |
| 3903 | self._num_threads = ops.convert_to_tensor( |
| 3904 | num_threads, dtype=dtypes.int64, name="num_threads") |
| 3905 | variant_tensor = ged_ops.private_thread_pool_dataset( |
| 3906 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3907 | self._num_threads, |
| 3908 | **self._flat_structure) |
| 3909 | super(_PrivateThreadPoolDataset, self).__init__(input_dataset, |
| 3910 | variant_tensor) |
| 3911 | |
| 3912 | |
| 3913 | class _RestructuredDataset(UnaryDataset): |