A `Dataset` that acts as an identity, overriding intra-op parallelism.
| 3879 | |
| 3880 | |
| 3881 | class _MaxIntraOpParallelismDataset(UnaryUnchangedStructureDataset): |
| 3882 | """A `Dataset` that acts as an identity, overriding intra-op parallelism.""" |
| 3883 | |
| 3884 | def __init__(self, input_dataset, max_intra_op_parallelism): |
| 3885 | self._input_dataset = input_dataset |
| 3886 | self._max_intra_op_parallelism = ops.convert_to_tensor( |
| 3887 | max_intra_op_parallelism, |
| 3888 | dtype=dtypes.int64, |
| 3889 | name="max_intra_op_parallelism") |
| 3890 | variant_tensor = ged_ops.max_intra_op_parallelism_dataset( |
| 3891 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3892 | self._max_intra_op_parallelism, |
| 3893 | **self._flat_structure) |
| 3894 | super(_MaxIntraOpParallelismDataset, self).__init__(input_dataset, |
| 3895 | variant_tensor) |
| 3896 | |
| 3897 | |
| 3898 | class _PrivateThreadPoolDataset(UnaryUnchangedStructureDataset): |