1# Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""The Counter Dataset.""" 16from tensorflow.python import tf2 17from tensorflow.python.data.ops import dataset_ops 18from tensorflow.python.framework import dtypes 19from tensorflow.python.framework import ops 20from tensorflow.python.util.tf_export import tf_export 21 22 23@tf_export("data.experimental.Counter", v1=[]) 24def CounterV2(start=0, step=1, dtype=dtypes.int64): 25 """Creates a `Dataset` that counts from `start` in steps of size `step`. 26 27 Unlike `tf.data.Dataset.range` which will stop at some ending number, 28 `Counter` will produce elements indefinitely. 29 30 >>> dataset = tf.data.experimental.Counter().take(5) 31 >>> list(dataset.as_numpy_iterator()) 32 [0, 1, 2, 3, 4] 33 >>> dataset.element_spec 34 TensorSpec(shape=(), dtype=tf.int64, name=None) 35 >>> dataset = tf.data.experimental.Counter(dtype=tf.int32) 36 >>> dataset.element_spec 37 TensorSpec(shape=(), dtype=tf.int32, name=None) 38 >>> dataset = tf.data.experimental.Counter(start=2).take(5) 39 >>> list(dataset.as_numpy_iterator()) 40 [2, 3, 4, 5, 6] 41 >>> dataset = tf.data.experimental.Counter(start=2, step=5).take(5) 42 >>> list(dataset.as_numpy_iterator()) 43 [2, 7, 12, 17, 22] 44 >>> dataset = tf.data.experimental.Counter(start=10, step=-1).take(5) 45 >>> list(dataset.as_numpy_iterator()) 46 [10, 9, 8, 7, 6] 47 48 Args: 49 start: (Optional.) The starting value for the counter. Defaults to 0. 50 step: (Optional.) The step size for the counter. Defaults to 1. 51 dtype: (Optional.) The data type for counter elements. Defaults to 52 `tf.int64`. 53 54 Returns: 55 A `Dataset` of scalar `dtype` elements. 56 """ 57 with ops.name_scope("counter"): 58 start = ops.convert_to_tensor(start, dtype=dtype, name="start") 59 step = ops.convert_to_tensor(step, dtype=dtype, name="step") 60 return dataset_ops.Dataset.from_tensors(0).repeat(None).scan( 61 start, lambda state, _: (state + step, state)) 62 63 64@tf_export(v1=["data.experimental.Counter"]) 65def CounterV1(start=0, step=1, dtype=dtypes.int64): 66 return dataset_ops.DatasetV1Adapter(CounterV2(start, step, dtype)) 67 68 69CounterV1.__doc__ = CounterV2.__doc__ 70 71if tf2.enabled(): 72 Counter = CounterV2 73else: 74 Counter = CounterV1 75