# Copyright 2017 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""Python wrappers for Datasets."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import abc

import enum

import functools

import threading

import warnings

import weakref

import numpy as np

import six

from six.moves import queue as Queue # pylint: disable=redefined-builtin

from tensorflow.core.framework import graph_pb2

from tensorflow.python import tf2

from tensorflow.python.compat import compat

from tensorflow.python.data.experimental.ops import distribute_options

from tensorflow.python.data.experimental.ops import optimization_options

from tensorflow.python.data.experimental.ops import stats_options

from tensorflow.python.data.experimental.ops import threading_options

from tensorflow.python.data.ops import iterator_ops

from tensorflow.python.data.util import nest

from tensorflow.python.data.util import options as options_lib

from tensorflow.python.data.util import random_seed

from tensorflow.python.data.util import sparse

from tensorflow.python.data.util import structure

from tensorflow.python.data.util import traverse

from tensorflow.python.eager import context

from tensorflow.python.eager import function as eager_function

from tensorflow.python.framework import composite_tensor

from tensorflow.python.framework import constant_op

from tensorflow.python.framework import dtypes

from tensorflow.python.framework import function

from tensorflow.python.framework import ops

from tensorflow.python.framework import random_seed as core_random_seed

from tensorflow.python.framework import smart_cond

from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib

from tensorflow.python.framework import tensor_shape

from tensorflow.python.framework import tensor_spec

from tensorflow.python.framework import tensor_util

from tensorflow.python.framework import type_spec

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import control_flow_ops

from tensorflow.python.ops import gen_dataset_ops

from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops

from tensorflow.python.ops import gen_io_ops

from tensorflow.python.ops import math_ops

from tensorflow.python.ops import script_ops

from tensorflow.python.ops import string_ops

from tensorflow.python.platform import tf_logging as logging

from tensorflow.python.training.tracking import base as tracking_base

from tensorflow.python.training.tracking import tracking

from tensorflow.python.util import deprecation

from tensorflow.python.util import function_utils

from tensorflow.python.util import lazy_loader

from tensorflow.python.util import nest as tf_nest

from tensorflow.python.util.tf_export import tf_export

# Loaded lazily due to a circular dependency (roughly

# tf.function->wrap_function->dataset->autograph->tf.function).

# TODO(b/133251390): Use a regular import.

wrap_function = lazy_loader.LazyLoader(

"wrap_function", globals(),

"tensorflow.python.eager.wrap_function")

# TODO(mdan): Create a public API for this.

autograph_ctx = lazy_loader.LazyLoader(

"autograph_ctx", globals(),

"tensorflow.python.autograph.core.ag_ctx")

autograph = lazy_loader.LazyLoader(

"autograph", globals(),

"tensorflow.python.autograph.impl.api")

ops.NotDifferentiable("ReduceDataset")

# A constant that can be used to enable auto-tuning.

AUTOTUNE = -1

tf_export("data.experimental.AUTOTUNE").export_constant(__name__, "AUTOTUNE")

class AutotuneAlgorithm(enum.Enum):

HILL_CLIMB = 0

GRADIENT_DESCENT = 1

@tf_export("data.Dataset", v1=[])

@six.add_metaclass(abc.ABCMeta)

class DatasetV2(tracking_base.Trackable, composite_tensor.CompositeTensor):

"""Represents a potentially large set of elements.

A `Dataset` can be used to represent an input pipeline as a

collection of elements and a "logical plan" of transformations that act on

those elements.

A dataset contains elements that each have the same (nested) structure and the

individual components of the structure can be of any type representable by

`tf.TypeSpec`, including `tf.Tensor`, `tf.data.Dataset`, `tf.SparseTensor`,

`tf.RaggedTensor`, or `tf.TensorArray`.

Example elements:

```python

# Integer element

a = 1

# Float element

b = 2.0

# Tuple element with 2 components

c = (1, 2)

# Dict element with 3 components

d = {"a": (2, 2), "b": 3}

# Element containing a dataset

e = tf.data.Dataset.from_element(10)

```

"""

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

weak_self = weakref.proxy(self)

self._variant_tracker = self._track_trackable(

_VariantTracker(

self._variant_tensor,

# _trace_variant_creation only works when executing eagerly, so we

# don't want to run it immediately. We also want the _VariantTracker

# to have a weak reference to the Dataset to avoid creating

# reference cycles and making work for the garbage collector.

lambda: weak_self._trace_variant_creation()()), # pylint: disable=unnecessary-lambda,protected-access

name="_variant_tracker")

self._graph_attr = ops.get_default_graph()

@property

def _variant_tensor(self):

return self._variant_tensor_attr

@_variant_tensor.setter

def _variant_tensor(self, _):

raise ValueError("The _variant_tensor property is read-only")

def _as_serialized_graph(self, stateful_whitelist=None):

"""Produces serialized graph representation of the dataset.

Args:

stateful_whitelist: Comma separated list of ops whose stateful attribute

should be ignored during serialization.

Returns:

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

serialized graph.

"""

if compat.forward_compatible(2019, 9, 10) or stateful_whitelist:

return gen_dataset_ops.dataset_to_graph(self._variant_tensor,

stateful_whitelist)

else:

return gen_dataset_ops.dataset_to_graph(self._variant_tensor)

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(

"Can only export Datasets which were created executing eagerly. "

"Please file a feature request if this is important to you.")

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

graph_def = graph_pb2.GraphDef().FromString(

self._as_serialized_graph().numpy()) # pylint: disable=protected-access

output_node_name = None

for node in graph_def.node:

if node.op == "_Retval":

if output_node_name is not None:

raise AssertionError(

"Found multiple return values from the dataset's graph, expected "

"only one.")

output_node_name, = node.input

if output_node_name is None:

raise AssertionError("Could not find the dataset's output node.")

# 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")

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("Dataset._inputs")

@property

def _graph(self):

return self._graph_attr

@_graph.setter

def _graph(self, _):

raise ValueError("The _graph property is read-only")

def _has_captured_ref(self):

"""Whether this dataset uses a function that captures ref variables.

Returns:

A boolean, which if true indicates that the dataset or one of its inputs

uses a function that captures ref variables.

"""

if context.executing_eagerly():

# RefVariables are not supported in eager mode

return False

def is_tensor_or_parent_ref(tensor):

if tensor.dtype._is_ref_dtype: # pylint: disable=protected-access

return True

return any([is_tensor_or_parent_ref(x) for x in tensor.op.inputs])

for fn in self._functions():

if any([is_tensor_or_parent_ref(t) for t in fn.function.captured_inputs]):

return True

return any(

[input_dataset._has_captured_ref() for input_dataset in self._inputs()]) # pylint: disable=protected-access

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

# by this dataset and its inputs.

def _functions(self):

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

Returns:

A list of `StructuredFunctionWrapper` objects.

"""

return []

def options(self):

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

Returns:

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

"""

options = Options()

for input_dataset in self._inputs():

input_options = input_dataset.options()

if input_options is not None:

options = options.merge(input_options)

return options

def _apply_options(self):

"""Apply options, such as optimization configuration, to the dataset."""

dataset = self

options = self.options()

if options.experimental_threading is not None:

t_options = options.experimental_threading

if t_options.max_intra_op_parallelism is not None:

dataset = _MaxIntraOpParallelismDataset(

dataset, t_options.max_intra_op_parallelism)

if t_options.private_threadpool_size is not None:

dataset = _PrivateThreadPoolDataset(dataset,

t_options.private_threadpool_size)

# pylint: disable=protected-access

static_optimizations = options._static_optimizations()

static_optimization_configs = options._static_optimization_configs()

# pylint: enable=protected-access

if static_optimizations:

if self._has_captured_ref():

warnings.warn(

"tf.data static optimizations are not compatible with tf.Variable. "

"The following optimizations will be disabled: %s. To enable "

"optimizations, use resource variables instead by calling "

"`tf.enable_resource_variables()` at the start of the program." %

", ".join(static_optimizations))

else:

dataset = _OptimizeDataset(dataset, static_optimizations,

static_optimization_configs)

autotune = True

algorithm = AutotuneAlgorithm.HILL_CLIMB

cpu_budget = 0 # Indicates that all CPU cores should be used.

if options.experimental_optimization is not None:

if options.experimental_optimization.autotune is False: # pylint: disable=g-bool-id-comparison

autotune = False

if options.experimental_optimization.autotune_algorithm is not None:

algorithm = options.experimental_optimization.autotune_algorithm

if options.experimental_optimization.autotune_cpu_budget is not None:

cpu_budget = options.experimental_optimization.autotune_cpu_budget

if autotune:

dataset = _ModelDataset(dataset, algorithm, cpu_budget)

if options.experimental_stats and options.experimental_stats.aggregator: # pylint: disable=line-too-long

dataset = _SetStatsAggregatorDataset( # pylint: disable=protected-access

dataset, options.experimental_stats.aggregator,

options.experimental_stats.prefix,

options.experimental_stats.counter_prefix)

return dataset

def __iter__(self):

"""Creates an `Iterator` for enumerating the elements of this dataset.

The returned iterator implements the Python iterator protocol and therefore

can only be used in eager mode.

Returns:

An `Iterator` over the elements of this dataset.

Raises:

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

"""

if (context.executing_eagerly()

or ops.get_default_graph()._building_function): # pylint: disable=protected-access

return iterator_ops.IteratorV2(self)

else:

raise RuntimeError("__iter__() is only supported inside of tf.function "

"or when eager execution is enabled.")

@abc.abstractproperty

def element_spec(self):

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

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("Dataset.element_spec")

def __repr__(self):

output_shapes = nest.map_structure(str, get_legacy_output_shapes(self))

output_shapes = str(output_shapes).replace("'", "")

output_types = nest.map_structure(repr, get_legacy_output_types(self))

output_types = str(output_types).replace("'", "")

return ("<%s shapes: %s, types: %s>" % (type(self).__name__, output_shapes,

output_types))

@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 _type_spec(self):

return DatasetSpec(self.element_spec)

@staticmethod

def from_tensors(tensors):

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

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/datasets#consuming_numpy_arrays).

Args:

tensors: A dataset element.

Returns:

Dataset: A `Dataset`.

"""

return TensorDataset(tensors)

@staticmethod

def from_tensor_slices(tensors):

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

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/datasets#consuming_numpy_arrays).

Args:

tensors: A dataset element, with each component having the same size in

the 0th dimension.

Returns:

Dataset: A `Dataset`.

"""

return TensorSliceDataset(tensors)

class _GeneratorState(object):

"""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 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):

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[iterator_id]

@staticmethod

def from_generator(generator, output_types, output_shapes=None, args=None):

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

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 the given

`output_types` and (optional) `output_shapes` arguments.

For example:

```python

import itertools

tf.compat.v1.enable_eager_execution()

def gen():

for i in itertools.count(1):

yield (i, [1] * i)

ds = tf.data.Dataset.from_generator(

gen, (tf.int64, tf.int64), (tf.TensorShape([]), tf.TensorShape([None])))

for value in ds.take(2):

print value

# (1, array([1]))

# (2, array([1, 1]))

```

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()`. 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.

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()`.

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: 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.

Returns:

Dataset: A `Dataset`.

"""

if not callable(generator):

raise TypeError("`generator` must be callable.")

if output_shapes is None:

output_shapes = nest.map_structure(

lambda _: tensor_shape.TensorShape(None), output_types)

else:

output_shapes = nest.map_structure_up_to(

output_types, tensor_shape.as_shape, output_shapes)

if args is None:

args = ()

else:

args = tuple(ops.convert_n_to_tensor(args, name="args"))

flattened_types = [dtypes.as_dtype(dt) for dt in nest.flatten(output_types)]

flattened_shapes = nest.flatten(output_shapes)

generator_state = DatasetV2._GeneratorState(generator)

def get_iterator_id_fn(unused_dummy):

"""Creates a unique `iterator_id` for each pass over the dataset.

The returned `iterator_id` disambiguates between multiple concurrently

existing iterators.

Args:

unused_dummy: Ignored value.

Returns:

A `tf.int64` tensor whose value uniquely identifies an iterator in

`generator_state`.

"""

return script_ops.numpy_function(generator_state.get_next_id, args,

dtypes.int64)

def generator_next_fn(iterator_id_t):

"""Generates the next element from iterator with ID `iterator_id_t`.

We map this function across an infinite repetition of the

`iterator_id_t`, and raise `StopIteration` to terminate the iteration.

Args:

iterator_id_t: A `tf.int64` tensor whose value uniquely identifies the

iterator in `generator_state` from which to generate an element.

Returns:

The next element to generate from the iterator.

"""

def generator_py_func(iterator_id):

"""A `py_func` that will be called to invoke the iterator."""

# `next()` raises `StopIteration` when there are no more

# elements remaining to be generated.

values = next(generator_state.get_iterator(iterator_id))

# Use the same _convert function from the py_func() implementation to

# convert the returned values to arrays early, so that we can inspect

# their values.

try:

flattened_values = nest.flatten_up_to(output_types, values)

except (TypeError, ValueError):

raise TypeError(

"`generator` yielded an element that did not match the expected "

"structure. The expected structure was %s, but the yielded "

"element was %s." % (output_types, values))

ret_arrays = []

for ret, dtype in zip(flattened_values, flattened_types):

try:

ret_arrays.append(script_ops.FuncRegistry._convert( # pylint: disable=protected-access

ret, dtype=dtype.as_numpy_dtype))

except (TypeError, ValueError):

raise TypeError(

"`generator` yielded an element that could not be converted to "

"the expected type. The expected type was %s, but the yielded "

"element was %s." % (dtype.name, ret))

# Additional type and shape checking to ensure that the components

# of the generated element match the `output_types` and `output_shapes`

# arguments.

for (ret_array, expected_dtype, expected_shape) in zip(

ret_arrays, flattened_types, flattened_shapes):

if ret_array.dtype != expected_dtype.as_numpy_dtype:

raise TypeError(

"`generator` yielded an element of type %s where an element "

"of type %s was expected." % (ret_array.dtype,

expected_dtype.as_numpy_dtype))

if not expected_shape.is_compatible_with(ret_array.shape):

raise ValueError(

"`generator` yielded an element of shape %s where an element "

"of shape %s was expected." % (ret_array.shape, expected_shape))

return ret_arrays

flat_values = script_ops.numpy_function(generator_py_func,

[iterator_id_t], flattened_types)

# The `py_func()` op drops the inferred shapes, so we add them back in

# here.

if output_shapes is not None:

for ret_t, shape in zip(flat_values, flattened_shapes):

ret_t.set_shape(shape)

return nest.pack_sequence_as(output_types, flat_values)

def finalize_fn(iterator_id_t):

"""Releases host-side state for the iterator with ID `iterator_id_t`."""

def finalize_py_func(iterator_id):

generator_state.iterator_completed(iterator_id)

# We return a dummy value so that the `finalize_fn` has a valid

# signature.

# 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(0, dtype=np.int64)

return script_ops.numpy_function(finalize_py_func, [iterator_id_t],

dtypes.int64)

# This function associates each traversal of `generator` with a unique

# iterator ID.

def flat_map_fn(dummy_arg):

# The `get_iterator_id_fn` gets a unique ID for the current instance of

# of the generator.

# The `generator_next_fn` gets the next element from the iterator with the

# given ID, and raises StopIteration when that iterator contains no

# more elements.

return _GeneratorDataset(dummy_arg, get_iterator_id_fn, generator_next_fn,

finalize_fn)

# A single-element dataset that, each time it is evaluated, contains a

# freshly-generated and unique (for the returned dataset) int64

# ID that will be used to identify the appropriate Python state, which

# is encapsulated in `generator_state`, and captured in

# `get_iterator_id_map_fn`.

dummy = 0

id_dataset = Dataset.from_tensors(dummy)

# A dataset that contains all of the elements generated by a

# single iterator created from `generator`, identified by the

# iterator ID contained in `id_dataset`. Lifting the iteration

# into a flat_map here enables multiple repetitions and/or nested

# versions of the returned dataset to be created, because it forces

# the generation of a new ID for each version.

return id_dataset.flat_map(flat_map_fn)

@staticmethod

def range(*args):

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

For example:

```python

Dataset.range(5) == [0, 1, 2, 3, 4]

Dataset.range(2, 5) == [2, 3, 4]

Dataset.range(1, 5, 2) == [1, 3]

Dataset.range(1, 5, -2) == []

Dataset.range(5, 1) == []

Dataset.range(5, 1, -2) == [5, 3]

```

Args:

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

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, stop = args[2]

Returns:

Dataset: A `RangeDataset`.

Raises:

ValueError: if len(args) == 0.

"""

return RangeDataset(*args)

@staticmethod

def zip(datasets):

"""Creates a `Dataset` by zipping together the given datasets.

This method has similar semantics to the built-in `zip()` function

in Python, with the main difference being that the `datasets`

argument can be an arbitrary nested structure of `Dataset` objects.

For example:

```python

a = Dataset.range(1, 4) # ==> [ 1, 2, 3 ]

b = Dataset.range(4, 7) # ==> [ 4, 5, 6 ]

c = Dataset.range(7, 13).batch(2) # ==> [ [7, 8], [9, 10], [11, 12] ]

d = Dataset.range(13, 15) # ==> [ 13, 14 ]

# The nested structure of the `datasets` argument determines the

# structure of elements in the resulting dataset.

Dataset.zip((a, b)) # ==> [ (1, 4), (2, 5), (3, 6) ]

Dataset.zip((b, a)) # ==> [ (4, 1), (5, 2), (6, 3) ]

# The `datasets` argument may contain an arbitrary number of

# datasets.

Dataset.zip((a, b, c)) # ==> [ (1, 4, [7, 8]),

# (2, 5, [9, 10]),

# (3, 6, [11, 12]) ]

# The number of elements in the resulting dataset is the same as

# the size of the smallest dataset in `datasets`.

Dataset.zip((a, d)) # ==> [ (1, 13), (2, 14) ]

```

Args:

datasets: A nested structure of datasets.

Returns:

Dataset: A `Dataset`.

"""

return ZipDataset(datasets)

def concatenate(self, dataset):

"""Creates a `Dataset` by concatenating the given dataset with this dataset.

```python

a = Dataset.range(1, 4) # ==> [ 1, 2, 3 ]

b = Dataset.range(4, 8) # ==> [ 4, 5, 6, 7 ]

# The input dataset and dataset to be concatenated should have the same

# nested structures and output types.

# c = Dataset.range(8, 14).batch(2) # ==> [ [8, 9], [10, 11], [12, 13] ]

# d = Dataset.from_tensor_slices([14.0, 15.0, 16.0])

# a.concatenate(c) and a.concatenate(d) would result in error.

a.concatenate(b) # ==> [ 1, 2, 3, 4, 5, 6, 7 ]

```

Args:

dataset: `Dataset` to be concatenated.

Returns:

Dataset: A `Dataset`.

"""

return ConcatenateDataset(self, dataset)

def prefetch(self, buffer_size):

"""Creates a `Dataset` that prefetches elements from this dataset.

Note: Like other `Dataset` methods, prefetch operates on the

elements of the input dataset. It has no concept of examples vs. batches.

`examples.prefetch(2)` will prefetch two elements (2 examples),

while `examples.batch(20).prefetch(2)` will prefetch 2 elements

(2 batches, of 20 examples each).

Args:

buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the maximum

number of elements that will be buffered when prefetching.

Returns:

Dataset: A `Dataset`.

"""

return PrefetchDataset(self, buffer_size)

@staticmethod

def list_files(file_pattern, shuffle=None, seed=None):

"""A dataset of all files matching one or more glob patterns.

NOTE: The default behavior of this method is to return filenames in

a non-deterministic random shuffled order. Pass a `seed` or `shuffle=False`

to get results in a deterministic order.

Example:

If we had the following files on our filesystem:

- /path/to/dir/a.txt

- /path/to/dir/b.py

- /path/to/dir/c.py

If we pass "/path/to/dir/*.py" as the directory, the dataset

would produce:

- /path/to/dir/b.py

- /path/to/dir/c.py

Args:

file_pattern: A string, a list of strings, or a `tf.Tensor` of string type

(scalar or vector), representing the filename glob (i.e. shell wildcard)

pattern(s) that will be matched.

shuffle: (Optional.) If `True`, the file names will be shuffled randomly.

Defaults to `True`.

seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the random

seed that will be used to create the distribution. See

`tf.compat.v1.set_random_seed` for behavior.

Returns:

Dataset: A `Dataset` of strings corresponding to file names.

"""

with ops.name_scope("list_files"):

if shuffle is None:

shuffle = True

file_pattern = ops.convert_to_tensor(

file_pattern, dtype=dtypes.string, name="file_pattern")

matching_files = gen_io_ops.matching_files(file_pattern)

# Raise an exception if `file_pattern` does not match any files.

condition = math_ops.greater(array_ops.shape(matching_files)[0], 0,

name="match_not_empty")

message = math_ops.add(

"No files matched pattern: ",

string_ops.reduce_join(file_pattern, separator=", "), name="message")

assert_not_empty = control_flow_ops.Assert(

condition, [message], summarize=1, name="assert_not_empty")

with ops.control_dependencies([assert_not_empty]):

matching_files = array_ops.identity(matching_files)

dataset = Dataset.from_tensor_slices(matching_files)

if shuffle:

# NOTE(mrry): The shuffle buffer size must be greater than zero, but the

# list of files might be empty.

buffer_size = math_ops.maximum(

array_ops.shape(matching_files, out_type=dtypes.int64)[0], 1)

dataset = dataset.shuffle(buffer_size, seed=seed)

return dataset

def repeat(self, count=None):

"""Repeats this dataset `count` times.

NOTE: If this dataset is a function of global state (e.g. a random number

generator), then different repetitions may produce different elements.

Args:

count: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the

number of times the dataset should be repeated. The default behavior (if

`count` is `None` or `-1`) is for the dataset be repeated indefinitely.

Returns:

Dataset: A `Dataset`.

"""

return RepeatDataset(self, count)

def enumerate(self, start=0):

"""Enumerates the elements of this dataset.

It is similar to python's `enumerate`.

For example:

```python

# NOTE: The following examples use `{ ... }` to represent the

# contents of a dataset.

a = { 1, 2, 3 }

b = { (7, 8), (9, 10) }

# The nested structure of the `datasets` argument determines the

# structure of elements in the resulting dataset.

a.enumerate(start=5)) == { (5, 1), (6, 2), (7, 3) }

b.enumerate() == { (0, (7, 8)), (1, (9, 10)) }

```

Args:

start: A `tf.int64` scalar `tf.Tensor`, representing the start value for

enumeration.

Returns:

Dataset: A `Dataset`.

"""

max_value = np.iinfo(dtypes.int64.as_numpy_dtype).max

return Dataset.zip((Dataset.range(start, max_value), self))

def shuffle(self, buffer_size, seed=None, reshuffle_each_iteration=None):

"""Randomly shuffles the elements of this dataset.

This dataset fills a buffer with `buffer_size` elements, then randomly

samples elements from this buffer, replacing the selected elements with new

elements. For perfect shuffling, a buffer size greater than or equal to the

full size of the dataset is required.

For instance, if your dataset contains 10,000 elements but `buffer_size` is

set to 1,000, then `shuffle` will initially select a random element from

only the first 1,000 elements in the buffer. Once an element is selected,

its space in the buffer is replaced by the next (i.e. 1,001-st) element,

maintaining the 1,000 element buffer.

Args:

buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the number of

elements from this dataset from which the new dataset will sample.

seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the random

seed that will be used to create the distribution. See

`tf.compat.v1.set_random_seed` for behavior.

reshuffle_each_iteration: (Optional.) A boolean, which if true indicates

that the dataset should be pseudorandomly reshuffled each time it is

iterated over. (Defaults to `True`.)

Returns:

Dataset: A `Dataset`.

"""

return ShuffleDataset(self, buffer_size, seed, reshuffle_each_iteration)

def cache(self, filename=""):

"""Caches the elements in this dataset.

Args:

filename: A `tf.string` scalar `tf.Tensor`, representing the name of a

directory on the filesystem to use for caching elements in this Dataset.

If a filename is not provided, the dataset will be cached in memory.

Returns:

Dataset: A `Dataset`.

"""

return CacheDataset(self, filename)

def take(self, count):

"""Creates a `Dataset` with at most `count` elements from this dataset.

Args:

count: A `tf.int64` scalar `tf.Tensor`, representing the number of

elements of this dataset that should be taken to form the new dataset.

If `count` is -1, or if `count` is greater than the size of this

dataset, the new dataset will contain all elements of this dataset.

Returns:

Dataset: A `Dataset`.

"""

return TakeDataset(self, count)

def skip(self, count):

"""Creates a `Dataset` that skips `count` elements from this dataset.

Args:

count: A `tf.int64` scalar `tf.Tensor`, representing the number of

elements of this dataset that should be skipped to form the new dataset.

If `count` is greater than the size of this dataset, the new dataset

will contain no elements. If `count` is -1, skips the entire dataset.

Returns:

Dataset: A `Dataset`.

"""

return SkipDataset(self, count)

def shard(self, num_shards, index):

"""Creates a `Dataset` that includes only 1/`num_shards` of this dataset.

This dataset operator is very useful when running distributed training, as

it allows each worker to read a unique subset.

When reading a single input file, you can skip elements as follows:

```python

d = tf.data.TFRecordDataset(input_file)

d = d.shard(num_workers, worker_index)

d = d.repeat(num_epochs)

d = d.shuffle(shuffle_buffer_size)

d = d.map(parser_fn, num_parallel_calls=num_map_threads)

```

Important caveats:

- Be sure to shard before you use any randomizing operator (such as

shuffle).

- Generally it is best if the shard operator is used early in the dataset

pipeline. For example, when reading from a set of TFRecord files, shard

before converting the dataset to input samples. This avoids reading every

file on every worker. The following is an example of an efficient

sharding strategy within a complete pipeline:

```python

d = Dataset.list_files(pattern)