A `Dataset` of a step separated range of values.
| 3004 | |
| 3005 | |
| 3006 | class RangeDataset(DatasetSource): |
| 3007 | """A `Dataset` of a step separated range of values.""" |
| 3008 | |
| 3009 | def __init__(self, *args): |
| 3010 | """See `Dataset.range()` for details.""" |
| 3011 | self._parse_args(*args) |
| 3012 | self._structure = tensor_spec.TensorSpec([], dtypes.int64) |
| 3013 | variant_tensor = gen_dataset_ops.range_dataset( |
| 3014 | start=self._start, |
| 3015 | stop=self._stop, |
| 3016 | step=self._step, |
| 3017 | **self._flat_structure) |
| 3018 | super(RangeDataset, self).__init__(variant_tensor) |
| 3019 | |
| 3020 | def _parse_args(self, *args): |
| 3021 | """Parse arguments according to the same rules as the `range()` builtin.""" |
| 3022 | if len(args) == 1: |
| 3023 | self._start = self._build_tensor(0, "start") |
| 3024 | self._stop = self._build_tensor(args[0], "stop") |
| 3025 | self._step = self._build_tensor(1, "step") |
| 3026 | elif len(args) == 2: |
| 3027 | self._start = self._build_tensor(args[0], "start") |
| 3028 | self._stop = self._build_tensor(args[1], "stop") |
| 3029 | self._step = self._build_tensor(1, "step") |
| 3030 | elif len(args) == 3: |
| 3031 | self._start = self._build_tensor(args[0], "start") |
| 3032 | self._stop = self._build_tensor(args[1], "stop") |
| 3033 | self._step = self._build_tensor(args[2], "step") |
| 3034 | else: |
| 3035 | raise ValueError("Invalid arguments to RangeDataset: %s" % str(args)) |
| 3036 | |
| 3037 | def _build_tensor(self, int64_value, name): |
| 3038 | return ops.convert_to_tensor(int64_value, dtype=dtypes.int64, name=name) |
| 3039 | |
| 3040 | @property |
| 3041 | def element_spec(self): |
| 3042 | return self._structure |
| 3043 | |
| 3044 | |
| 3045 | class _MemoryCacheDeleter(object): |