if not name.isidentifier():

raise ValueError("Invalid `name`. The argument `name` needs to be a valid "

"identifier. Value is considered a valid identifier if it "

"only contains alphanumeric characters (a-z), (A-Z), and "

"(0-9), or underscores (_). A valid identifier cannot "

"start with a number, or contain any spaces.")

"""Returns the type of `value` if it is a TypeSpec."""

if isinstance(value, type_spec.TypeSpec):

return value.value_type()

else:

collections_abc.Iterable,

tracking_base.Trackable,

composite_tensor.CompositeTensor,

data_types.DatasetV2,

metaclass=abc.ABCMeta):

"""Represents a potentially large set of elements.

The `tf.data.Dataset` API supports writing descriptive and efficient input

pipelines. `Dataset` usage follows a common pattern:

1. Create a source dataset from your input data.

2. Apply dataset transformations to preprocess the data.

3. Iterate over the dataset and process the elements.

Iteration happens in a streaming fashion, so the full dataset does not need to

fit into memory.

Source Datasets:

The simplest way to create a dataset is to create it from a python `list`:

>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

>>> for element in dataset:

... print(element)

tf.Tensor(1, shape=(), dtype=int32)

tf.Tensor(2, shape=(), dtype=int32)

tf.Tensor(3, shape=(), dtype=int32)

To process lines from files, use `tf.data.TextLineDataset`:

>>> dataset = tf.data.TextLineDataset(["file1.txt", "file2.txt"])

To process records written in the `TFRecord` format, use `TFRecordDataset`:

>>> dataset = tf.data.TFRecordDataset(["file1.tfrecords", "file2.tfrecords"])

To create a dataset of all files matching a pattern, use

`tf.data.Dataset.list_files`:

```python

dataset = tf.data.Dataset.list_files("/path/*.txt")

```

See `tf.data.FixedLengthRecordDataset` and `tf.data.Dataset.from_generator`

for more ways to create datasets.

Transformations:

Once you have a dataset, you can apply transformations to prepare the data for

your model:

>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

>>> dataset = dataset.map(lambda x: x*2)

>>> list(dataset.as_numpy_iterator())

[2, 4, 6]

Common Terms:

**Element**: A single output from calling `next()` on a dataset iterator.

Elements may be nested structures containing multiple components. For

example, the element `(1, (3, "apple"))` has one tuple nested in another

tuple. The components are `1`, `3`, and `"apple"`.

**Component**: The leaf in the nested structure of an element.

Supported types:

Elements can be nested structures of tuples, named tuples, and dictionaries.

Note that Python lists are *not* treated as nested structures of components.

Instead, lists are converted to tensors and treated as components. For

example, the element `(1, [1, 2, 3])` has only two components; the tensor `1`

and the tensor `[1, 2, 3]`. Element components can be of any type

representable by `tf.TypeSpec`, including `tf.Tensor`, `tf.data.Dataset`,

`tf.sparse.SparseTensor`, `tf.RaggedTensor`, and `tf.TensorArray`.

```python

a = 1 # Integer element

b = 2.0 # Float element

c = (1, 2) # Tuple element with 2 components

d = {"a": (2, 2), "b": 3} # Dict element with 3 components

Point = collections.namedtuple("Point", ["x", "y"])

e = Point(1, 2) # Named tuple

f = tf.data.Dataset.range(10) # Dataset element

```

For more information,

read [this guide](https://www.tensorflow.org/guide/data).

"""

def __init__(self, variant_tensor):

"""Creates a DatasetV2 object.

This is a difference between DatasetV1 and DatasetV2. DatasetV1 does not

take anything in its constructor whereas in the DatasetV2, we expect

subclasses to create a variant_tensor and pass it in to the super() call.

Args:

variant_tensor: A DT_VARIANT tensor that represents the dataset.

"""

self._variant_tensor_attr = variant_tensor

self._graph_attr = ops.get_default_graph()

# Initialize the options for this dataset and its inputs.

self._options_attr = options_lib.Options()

for input_dataset in self._inputs():

input_options = None

if isinstance(input_dataset, data_types.DatasetV1):

# If the V1 dataset does not have the `_dataset` attribute, we assume it

# is a dataset source and hence does not have options. Otherwise, we

# grab the options of `_dataset` object

if hasattr(input_dataset, "_dataset"):

if not isinstance(input_dataset._dataset, data_types.DatasetV2):

raise TypeError(

f"Each input of dataset {type(self)} should be a subclass of "

f"`tf.data.Dataset` but encountered "

f"{type(input_dataset._dataset)}.")

input_options = input_dataset._dataset._options_attr

elif isinstance(input_dataset, data_types.DatasetV2):

input_options = input_dataset._options_attr

else:

raise TypeError(

f"Each input of dataset {type(self)} should be a subclass of "

f"`tf.data.Dataset` but encountered {type(input_dataset)}.")

if input_options is not None:

self._options_attr = self._options_attr.merge(input_options)

self._options_attr._set_mutable(False) # pylint: disable=protected-access

@property

def _variant_tensor(self):

return self._variant_tensor_attr

@_variant_tensor.setter

def _variant_tensor(self, _):

raise ValueError("The `_variant_tensor` property cannot be modified.")

@deprecation.deprecated_args(None, "Use external_state_policy instead",

"allow_stateful")

def _as_serialized_graph(

self,

allow_stateful=None,

strip_device_assignment=None,

external_state_policy=options_lib.ExternalStatePolicy.WARN):

"""Produces serialized graph representation of the dataset.

Args:

allow_stateful: If true, we allow stateful ops to be present in the graph

def. In that case, the state in these ops would be thrown away.

strip_device_assignment: If true, non-local (i.e. job and task) device

assignment is stripped from ops in the serialized graph.

external_state_policy: The ExternalStatePolicy enum that determines how we

handle input pipelines that depend on external state. By default, its

set to WARN.

Returns:

A scalar `tf.Tensor` of `tf.string` type, representing this dataset as a

serialized graph.

"""

if external_state_policy:

policy = external_state_policy.value

return gen_dataset_ops.dataset_to_graph_v2(

self._variant_tensor,

external_state_policy=policy,

strip_device_assignment=strip_device_assignment)

if strip_device_assignment:

return gen_dataset_ops.dataset_to_graph(

self._variant_tensor,

allow_stateful=allow_stateful,

strip_device_assignment=strip_device_assignment)

return gen_dataset_ops.dataset_to_graph(

self._variant_tensor, allow_stateful=allow_stateful)

def _maybe_track_assets(self, graph_def):

"""Finds and tracks nodes in `graph_def` that refer to asset files.

Args:

graph_def: Serialized graph representation of this dataset.

Returns:

A dictionary mapping the node name of an asset constant to a tracked

`asset.Asset` object.

"""

asset_tracker = {}

for node in graph_def.node:

if node.name.startswith("FileIdentity"):

asset_tracker[node.input[0]] = None

if not asset_tracker:

return {}

for node in graph_def.node:

if node.name in asset_tracker:

tensor_proto = node.attr["value"].tensor

with context.eager_mode(), ops.device("CPU"):

node_value = gen_parsing_ops.parse_tensor(

tensor_proto.SerializeToString(), dtypes.string).numpy()

asset_tracker[node.name] = ([

self._track_trackable(asset.Asset(n),

name=node.name + "_" + str(i), overwrite=True)

for i, n in enumerate(node_value)

])

return asset_tracker

def _trackable_children(self,

save_type=tracking_base.SaveType.CHECKPOINT,

**kwargs):

if save_type != tracking_base.SaveType.SAVEDMODEL:

return {}

# _trace_variant_creation only works when executing eagerly, so we don't

# want to run it in the object initialization.

@def_function.function(input_signature=[], autograph=False)

def _creator():

resource = self._trace_variant_creation()() # pylint: disable=protected-access

return resource

_creator.get_concrete_function() # Trigger asset tracking

children = super(DatasetV2, self)._trackable_children(save_type, **kwargs)

children["_variant_tracker"] = _VariantTracker(self._variant_tensor,

_creator)

return children

def _trace_variant_creation(self):

"""Traces a function which outputs a variant `tf.Tensor` for this dataset.

Note that creating this function involves evaluating an op, and is currently

only supported when executing eagerly.

Returns:

A zero-argument `ConcreteFunction` which outputs a variant `tf.Tensor`.

"""

variant = self._variant_tensor

if not isinstance(variant, ops.EagerTensor):

raise NotImplementedError(

"Constructing a tf.function that reproduces a given dataset is only "

"supported for datasets created eagerly. Please file a feature "

"request if this is important to you.")

with context.eager_mode(), ops.device("CPU"):

# pylint: disable=protected-access

graph_def = graph_pb2.GraphDef().FromString(

self._as_serialized_graph(external_state_policy=options_lib

.ExternalStatePolicy.FAIL).numpy())

output_node_names = []

for node in graph_def.node:

if node.op == "_Retval":

output_node_names = node.input

if len(output_node_names) != 1:

raise AssertionError(

f"Dataset graph is expected to only have one return value but found "

f"{len(output_node_names)} return values: {output_node_names}.")

output_node_name = output_node_names[0]

file_path_nodes = {}

# When building a tf.function, track files as `saved_model.Asset`s.

if ops.get_default_graph().building_function:

asset_tracker = self._maybe_track_assets(graph_def)

for key in asset_tracker:

assets_list = [

array_ops.expand_dims(asset.asset_path, axis=0)

for asset in asset_tracker[key]

]

file_path_nodes[key] = array_ops.concat(assets_list, axis=0)

# Add functions used in this Dataset to the function's graph, since they

# need to follow it around (and for example be added to a SavedModel which

# references the dataset).

variant_function = wrap_function.function_from_graph_def(

graph_def,

inputs=[],

outputs=output_node_name + ":0",

captures=file_path_nodes)

for used_function in self._functions():

used_function.function.add_to_graph(variant_function.graph)

return variant_function

@abc.abstractmethod

def _inputs(self):

"""Returns a list of the input datasets of the dataset."""

raise NotImplementedError(f"{type(self)}._inputs()")

@property

def _graph(self):

return self._graph_attr

@_graph.setter

def _graph(self, _):

raise ValueError("The `_graph` property cannot be modified.")

# TODO(jsimsa): Change this to be the transitive closure of functions used

# by this dataset and its inputs.

def _functions(self) -> list[StructuredFunctionWrapper]:

"""Returns a list of functions associated with this dataset.

Returns:

A list of `StructuredFunctionWrapper` objects.

"""

return []

def _options(self):

"""Returns the options tensor for this dataset."""

# pylint: disable=protected-access

return gen_dataset_ops.get_options(self._variant_tensor)

@classmethod

def _options_tensor_to_options(cls, serialized_options):

"""Converts options tensor to tf.data.Options object."""

options = options_lib.Options()

if tensor_util.constant_value(serialized_options) is not None:

pb = dataset_options_pb2.Options.FromString(tensor_util.constant_value(

serialized_options))

options._from_proto(pb) # pylint: disable=protected-access

return options

def options(self):

"""Returns the options for this dataset and its inputs.

Returns:

A `tf.data.Options` object representing the dataset options.

"""

if context.executing_eagerly():

options = self._options_tensor_to_options(self._options())

options._set_mutable(False) # pylint: disable=protected-access

return options

warnings.warn("To make it possible to preserve tf.data options across "

"serialization boundaries, their implementation has moved to "

"be part of the TensorFlow graph. As a consequence, the "

"options value is in general no longer known at graph "

"construction time. Invoking this method in graph mode "

"retains the legacy behavior of the original implementation, "

"but note that the returned value might not reflect the "

"actual value of the options.")

return self._options_attr

def _apply_debug_options(self):

if debug_mode.DEBUG_MODE:

# Disable autotuning and static optimizations that could introduce

# parallelism or asynchrony.

options = options_lib.Options()

options.autotune.enabled = False

options.experimental_optimization.filter_parallelization = False

options.experimental_optimization.map_and_batch_fusion = False

options.experimental_optimization.map_parallelization = False

dataset = _OptionsDataset(self, options)

else:

dataset = self

return dataset

def __iter__(self) -> iterator_ops.OwnedIterator:

"""Creates an iterator for elements of this dataset.

The returned iterator implements the Python Iterator protocol.

Returns:

An `tf.data.Iterator` for the elements of this dataset.

Raises:

RuntimeError: If not inside of tf.function and not executing eagerly.

"""

if context.executing_eagerly() or ops.inside_function():

with ops.colocate_with(self._variant_tensor):

return iterator_ops.OwnedIterator(self)

else:

raise RuntimeError("`tf.data.Dataset` only supports Python-style "

"iteration in eager mode or within tf.function.")

def __bool__(self):

return True # Required as __len__ is defined

__nonzero__ = __bool__ # Python 2 backward compatibility

def __len__(self):

"""Returns the length of the dataset if it is known and finite.

This method requires that you are running in eager mode, and that the

length of the dataset is known and non-infinite. When the length may be

unknown or infinite, or if you are running in graph mode, use

`tf.data.Dataset.cardinality` instead.

Returns:

An integer representing the length of the dataset.

Raises:

RuntimeError: If the dataset length is unknown or infinite, or if eager

execution is not enabled.

"""

if not context.executing_eagerly():

raise TypeError("`tf.data.Dataset` only supports `len` in eager mode. "

"Use `tf.data.Dataset.cardinality()` instead.")

length = self.cardinality()

if length.numpy() == INFINITE:

raise TypeError("The dataset is infinite.")

if length.numpy() == UNKNOWN:

raise TypeError("The dataset length is unknown.")

return length

@abc.abstractproperty

def element_spec(self):

"""The type specification of an element of this dataset.

>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

>>> dataset.element_spec

TensorSpec(shape=(), dtype=tf.int32, name=None)

For more information,

read [this guide](https://www.tensorflow.org/guide/data#dataset_structure).

Returns:

A (nested) structure of `tf.TypeSpec` objects matching the structure of an

element of this dataset and specifying the type of individual components.

"""

raise NotImplementedError(f"{type(self)}.element_spec()")

def __repr__(self):

type_ = type(self._dataset if isinstance(self, DatasetV1Adapter) else self)

return f"<{type_.__name__} element_spec={self.element_spec}>"

def __debug_string__(self):

"""Returns a string showing the type of the dataset and its inputs.

This string is intended only for debugging purposes, and may change without

warning.

"""

lines = []

to_process = [(self, 0)] # Stack of (dataset, depth) pairs.

while to_process:

dataset, depth = to_process.pop()

lines.append("-"*2*depth + repr(dataset))

to_process.extend([(ds, depth+1) for ds in dataset._inputs()]) # pylint: disable=protected-access

return "\n".join(lines)

def as_numpy_iterator(self):

"""Returns an iterator which converts all elements of the dataset to numpy.

Use `as_numpy_iterator` to inspect the content of your dataset. To see

element shapes and types, print dataset elements directly instead of using

`as_numpy_iterator`.

>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

>>> for element in dataset:

... print(element)

tf.Tensor(1, shape=(), dtype=int32)

tf.Tensor(2, shape=(), dtype=int32)

tf.Tensor(3, shape=(), dtype=int32)

This method requires that you are running in eager mode and the dataset's

element_spec contains only `TensorSpec` components.

>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

>>> for element in dataset.as_numpy_iterator():

... print(element)

1

2

3

>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

>>> print(list(dataset.as_numpy_iterator()))

[1, 2, 3]

`as_numpy_iterator()` will preserve the nested structure of dataset

elements.

>>> dataset = tf.data.Dataset.from_tensor_slices({'a': ([1, 2], [3, 4]),

... 'b': [5, 6]})

>>> list(dataset.as_numpy_iterator()) == [{'a': (1, 3), 'b': 5},

... {'a': (2, 4), 'b': 6}]

True

Returns:

An iterable over the elements of the dataset, with their tensors converted

to numpy arrays.

Raises:

TypeError: if an element contains a non-`Tensor` value.

RuntimeError: if eager execution is not enabled.

"""

if not context.executing_eagerly():

raise RuntimeError("`tf.data.Dataset.as_numpy_iterator()` is only "

"supported in eager mode.")

for component_spec in nest.flatten(self.element_spec):

if not isinstance(

component_spec,

(tensor_spec.TensorSpec, ragged_tensor.RaggedTensorSpec,

sparse_tensor_lib.SparseTensorSpec, none_tensor.NoneTensorSpec)):

raise TypeError(

f"`tf.data.Dataset.as_numpy_iterator()` is not supported for "

f"datasets that produce values of type {component_spec.value_type}")

return NumpyIterator(self)

@property

def _flat_shapes(self):

"""Returns a list `tf.TensorShapes`s for the element tensor representation.

Returns:

A list `tf.TensorShapes`s for the element tensor representation.

"""

return structure.get_flat_tensor_shapes(self.element_spec)

@property

def _flat_types(self):

"""Returns a list `tf.DType`s for the element tensor representation.

Returns:

A list `tf.DType`s for the element tensor representation.

"""

return structure.get_flat_tensor_types(self.element_spec)

@property

def _flat_structure(self):

"""Helper for setting `output_shapes` and `output_types` attrs of an op.

Most dataset op constructors expect `output_shapes` and `output_types`

arguments that represent the flattened structure of an element. This helper

function generates these attrs as a keyword argument dictionary, allowing

`Dataset._variant_tensor` implementations to pass `**self._flat_structure`

to the op constructor.

Returns:

A dictionary of keyword arguments that can be passed to a dataset op

constructor.

"""

return {

"output_shapes": self._flat_shapes,

"output_types": self._flat_types,

}

@property

def _metadata(self):

"""Helper for generating dataset metadata."""

metadata = dataset_metadata_pb2.Metadata()

if self._name:

metadata.name = _validate_and_encode(self._name)

return metadata

@property

def _common_args(self):

"""Helper for generating arguments that are common across most dataset ops.

Most dataset op constructors expect `output_shapes` and `output_types`

arguments that represent the flattened structure of an element, as well as a

`metadata` argument for additional metadata such as user-defined dataset

name. This helper function generates common attributes as a keyword argument

dictionary, allowing `Dataset._variant_tensor` implementations to pass

`**self._common_args` to the op constructor.

Returns:

A dictionary of keyword arguments that can be passed to a dataset op

constructor.

"""

return {

"metadata": self._metadata.SerializeToString(),

"output_shapes": self._flat_shapes,

"output_types": self._flat_types,

}

@property

def _type_spec(self):

return DatasetSpec(self.element_spec)

@staticmethod

def from_tensors(tensors, name=None) -> "DatasetV2":

"""Creates a `Dataset` with a single element, comprising the given tensors.

`from_tensors` produces a dataset containing only a single element. To slice

the input tensor into multiple elements, use `from_tensor_slices` instead.

>>> dataset = tf.data.Dataset.from_tensors([1, 2, 3])

>>> list(dataset.as_numpy_iterator())

[array([1, 2, 3], dtype=int32)]

>>> dataset = tf.data.Dataset.from_tensors(([1, 2, 3], 'A'))

>>> list(dataset.as_numpy_iterator())

[(array([1, 2, 3], dtype=int32), b'A')]

>>> # You can use `from_tensors` to produce a dataset which repeats

>>> # the same example many times.

>>> example = tf.constant([1,2,3])

>>> dataset = tf.data.Dataset.from_tensors(example).repeat(2)

>>> list(dataset.as_numpy_iterator())

[array([1, 2, 3], dtype=int32), array([1, 2, 3], dtype=int32)]

Note that if `tensors` contains a NumPy array, and eager execution is not

enabled, the values will be embedded in the graph as one or more

`tf.constant` operations. For large datasets (> 1 GB), this can waste

memory and run into byte limits of graph serialization. If `tensors`

contains one or more large NumPy arrays, consider the alternative described

in [this

guide](https://tensorflow.org/guide/data#consuming_numpy_arrays).

Args:

tensors: A dataset "element". Supported values are documented

[here](https://www.tensorflow.org/guide/data#dataset_structure).

name: (Optional.) A name for the tf.data operation.

Returns:

Dataset: A `Dataset`.

"""

# Loaded lazily due to a circular dependency (dataset_ops ->

# from_tensors_op -> dataset_ops).

# pylint: disable=g-import-not-at-top,protected-access

from tensorflow.python.data.ops import from_tensors_op

return from_tensors_op._from_tensors(tensors, name)

# pylint: enable=g-import-not-at-top,protected-access

@staticmethod

def from_tensor_slices(tensors, name=None) -> "DatasetV2":

"""Creates a `Dataset` whose elements are slices of the given tensors.

The given tensors are sliced along their first dimension. This operation

preserves the structure of the input tensors, removing the first dimension

of each tensor and using it as the dataset dimension. All input tensors

must have the same size in their first dimensions.

>>> # Slicing a 1D tensor produces scalar tensor elements.

>>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

>>> list(dataset.as_numpy_iterator())

[1, 2, 3]

>>> # Slicing a 2D tensor produces 1D tensor elements.

>>> dataset = tf.data.Dataset.from_tensor_slices([[1, 2], [3, 4]])

>>> list(dataset.as_numpy_iterator())

[array([1, 2], dtype=int32), array([3, 4], dtype=int32)]

>>> # Slicing a tuple of 1D tensors produces tuple elements containing

>>> # scalar tensors.

>>> dataset = tf.data.Dataset.from_tensor_slices(([1, 2], [3, 4], [5, 6]))

>>> list(dataset.as_numpy_iterator())

[(1, 3, 5), (2, 4, 6)]

>>> # Dictionary structure is also preserved.

>>> dataset = tf.data.Dataset.from_tensor_slices({"a": [1, 2], "b": [3, 4]})

>>> list(dataset.as_numpy_iterator()) == [{'a': 1, 'b': 3},

... {'a': 2, 'b': 4}]

True

>>> # Two tensors can be combined into one Dataset object.

>>> features = tf.constant([[1, 3], [2, 1], [3, 3]]) # ==> 3x2 tensor

>>> labels = tf.constant(['A', 'B', 'A']) # ==> 3x1 tensor

>>> dataset = Dataset.from_tensor_slices((features, labels))

>>> # Both the features and the labels tensors can be converted

>>> # to a Dataset object separately and combined after.

>>> features_dataset = Dataset.from_tensor_slices(features)

>>> labels_dataset = Dataset.from_tensor_slices(labels)

>>> dataset = Dataset.zip((features_dataset, labels_dataset))

>>> # A batched feature and label set can be converted to a Dataset

>>> # in similar fashion.

>>> batched_features = tf.constant([[[1, 3], [2, 3]],

... [[2, 1], [1, 2]],

... [[3, 3], [3, 2]]], shape=(3, 2, 2))

>>> batched_labels = tf.constant([['A', 'A'],

... ['B', 'B'],

... ['A', 'B']], shape=(3, 2, 1))

>>> dataset = Dataset.from_tensor_slices((batched_features, batched_labels))

>>> for element in dataset.as_numpy_iterator():

... print(element)

(array([[1, 3],

[2, 3]], dtype=int32), array([[b'A'],

[b'A']], dtype=object))

(array([[2, 1],

[1, 2]], dtype=int32), array([[b'B'],

[b'B']], dtype=object))

(array([[3, 3],

[3, 2]], dtype=int32), array([[b'A'],

[b'B']], dtype=object))

Note that if `tensors` contains a NumPy array, and eager execution is not

enabled, the values will be embedded in the graph as one or more

`tf.constant` operations. For large datasets (> 1 GB), this can waste

memory and run into byte limits of graph serialization. If `tensors`

contains one or more large NumPy arrays, consider the alternative described

in [this guide](

https://tensorflow.org/guide/data#consuming_numpy_arrays).

Args:

tensors: A dataset element, whose components have the same first

dimension. Supported values are documented

[here](https://www.tensorflow.org/guide/data#dataset_structure).

name: (Optional.) A name for the tf.data operation.

Returns:

Dataset: A `Dataset`.

"""

# Loaded lazily due to a circular dependency (dataset_ops ->

# from_tensor_slices_op -> dataset_ops).

# pylint: disable=g-import-not-at-top,protected-access

from tensorflow.python.data.ops import from_tensor_slices_op

return from_tensor_slices_op._from_tensor_slices(tensors, name)

# pylint: enable=g-import-not-at-top,protected-access

class _GeneratorState:

"""Stores outstanding iterators created from a Python generator.

This class keeps track of potentially multiple iterators that may have

been created from a generator, e.g. in the case that the dataset is

repeated, or nested within a parallel computation.

"""

def __init__(self, generator):

self._generator = generator

self._lock = threading.Lock()

self._next_id = 0 # GUARDED_BY(self._lock)

self._args = {}

self._iterators = {}

def _normalize_id(self, iterator_id):

# In debug mode, iterator ids may be eagerly-generated np.arrays instead

# of Tensors. We convert them to scalars to make them hashable.

if isinstance(iterator_id, np.ndarray):

return iterator_id.item()

return iterator_id

def get_next_id(self, *args):

with self._lock:

ret = self._next_id

self._next_id += 1

self._args[ret] = args

# NOTE(mrry): Explicitly create an array of `np.int64` because implicit

# casting in `py_func()` will create an array of `np.int32` on Windows,

# leading to a runtime error.

return np.array(ret, dtype=np.int64)

def get_iterator(self, iterator_id):

iterator_id = self._normalize_id(iterator_id)

try:

return self._iterators[iterator_id]

except KeyError:

iterator = iter(self._generator(*self._args.pop(iterator_id)))

self._iterators[iterator_id] = iterator

return iterator

def iterator_completed(self, iterator_id):

del self._iterators[self._normalize_id(iterator_id)]

@staticmethod

@deprecation.deprecated_args(None, "Use output_signature instead",

"output_types", "output_shapes")

def from_generator(

generator,

output_types=None,

output_shapes=None,

args=None,

output_signature=None,

name=None,

) -> "DatasetV2":

"""Creates a `Dataset` whose elements are generated by `generator`.

Note: The current implementation of `Dataset.from_generator()` uses

`tf.numpy_function` and inherits the same constraints. In particular, it

requires the dataset and iterator related operations to be placed

on a device in the same process as the Python program that called

`Dataset.from_generator()`. In particular, using `from_generator` will

preclude the use of tf.data service for scaling out dataset processing.

The body of `generator` will not be serialized in a `GraphDef`, and you

should not use this method if you need to serialize your model and restore

it in a different environment.

The `generator` argument must be a callable object that returns

an object that supports the `iter()` protocol (e.g. a generator function).

The elements generated by `generator` must be compatible with either the

given `output_signature` argument or with the given `output_types` and

(optionally) `output_shapes` arguments, whichever was specified.

The recommended way to call `from_generator` is to use the

`output_signature` argument. In this case the output will be assumed to

consist of objects with the classes, shapes and types defined by

`tf.TypeSpec` objects from `output_signature` argument:

>>> def gen():

... ragged_tensor = tf.ragged.constant([[1, 2], [3]])

... yield 42, ragged_tensor

>>>

>>> dataset = tf.data.Dataset.from_generator(

... gen,

... output_signature=(

... tf.TensorSpec(shape=(), dtype=tf.int32),

... tf.RaggedTensorSpec(shape=(2, None), dtype=tf.int32)))

>>>

>>> list(dataset.take(1))

[(<tf.Tensor: shape=(), dtype=int32, numpy=42>,

<tf.RaggedTensor [[1, 2], [3]]>)]

There is also a deprecated way to call `from_generator` by either with

`output_types` argument alone or together with `output_shapes` argument.

In this case the output of the function will be assumed to consist of

`tf.Tensor` objects with the types defined by `output_types` and with the

shapes which are either unknown or defined by `output_shapes`.

Note: If `generator` depends on mutable global variables or other external

state, be aware that the runtime may invoke `generator` multiple times

(in order to support repeating the `Dataset`) and at any time

between the call to `Dataset.from_generator()` and the production of the

first element from the generator. Mutating global variables or external

state can cause undefined behavior, and we recommend that you explicitly

cache any external state in `generator` before calling

`Dataset.from_generator()`.

Note: While the `output_signature` parameter makes it possible to yield

`Dataset` elements, the scope of `Dataset.from_generator()` should be

limited to logic that cannot be expressed through tf.data operations. Using

tf.data operations within the generator function is an anti-pattern and may

result in incremental memory growth.

Args:

generator: A callable object that returns an object that supports the

`iter()` protocol. If `args` is not specified, `generator` must take no

arguments; otherwise it must take as many arguments as there are values

in `args`.

output_types: (Optional.) A (nested) structure of `tf.DType` objects

corresponding to each component of an element yielded by `generator`.

output_shapes: (Optional.) A (nested) structure of `tf.TensorShape`

objects corresponding to each component of an element yielded by

`generator`.

args: (Optional.) A tuple of `tf.Tensor` objects that will be evaluated

and passed to `generator` as NumPy-array arguments.

output_signature: (Optional.) A (nested) structure of `tf.TypeSpec`

objects corresponding to each component of an element yielded by

`generator`.

name: (Optional.) A name for the tf.data operations used by

`from_generator`.

Returns:

Dataset: A `Dataset`.

"""

# Loaded lazily due to a circular dependency (dataset_ops ->

# from_generator_op -> dataset_ops).

# pylint: disable=g-import-not-at-top,protected-access

from tensorflow.python.data.ops import from_generator_op

return from_generator_op._from_generator(generator, output_types,

output_shapes, args,

output_signature, name)

# pylint: enable=g-import-not-at-top,protected-access

@staticmethod

def range(*args, **kwargs) -> "DatasetV2":

"""Creates a `Dataset` of a step-separated range of values.

>>> list(Dataset.range(5).as_numpy_iterator())

[0, 1, 2, 3, 4]

>>> list(Dataset.range(2, 5).as_numpy_iterator())

[2, 3, 4]

>>> list(Dataset.range(1, 5, 2).as_numpy_iterator())

[1, 3]

>>> list(Dataset.range(1, 5, -2).as_numpy_iterator())

[]

>>> list(Dataset.range(5, 1).as_numpy_iterator())

[]

>>> list(Dataset.range(5, 1, -2).as_numpy_iterator())

[5, 3]

>>> list(Dataset.range(2, 5, output_type=tf.int32).as_numpy_iterator())

[2, 3, 4]

>>> list(Dataset.range(1, 5, 2, output_type=tf.float32).as_numpy_iterator())

[1.0, 3.0]

Args:

*args: follows the same semantics as python's range.

len(args) == 1 -> start = 0, stop = args[0], step = 1.

len(args) == 2 -> start = args[0], stop = args[1], step = 1.

len(args) == 3 -> start = args[0], stop = args[1], step = args[2].

**kwargs:

- output_type: Its expected dtype. (Optional, default: `tf.int64`).