A `Dataset` for sharding its input.
| 3235 | |
| 3236 | |
| 3237 | class ShardDataset(UnaryUnchangedStructureDataset): |
| 3238 | """A `Dataset` for sharding its input.""" |
| 3239 | |
| 3240 | def __init__(self, input_dataset, num_shards, index): |
| 3241 | """See `Dataset.shard()` for details.""" |
| 3242 | self._input_dataset = input_dataset |
| 3243 | self._num_shards = ops.convert_to_tensor( |
| 3244 | num_shards, dtype=dtypes.int64, name="num_shards") |
| 3245 | self._index = ops.convert_to_tensor(index, dtype=dtypes.int64, name="index") |
| 3246 | variant_tensor = gen_dataset_ops.shard_dataset( |
| 3247 | input_dataset._variant_tensor, # pylint: disable=protected-access |
| 3248 | num_shards=self._num_shards, |
| 3249 | index=self._index, |
| 3250 | **self._flat_structure) |
| 3251 | super(ShardDataset, self).__init__(input_dataset, variant_tensor) |
| 3252 | |
| 3253 | |
| 3254 | class BatchDataset(UnaryDataset): |