A Dataset comprising lines from one or more text files.
Inherits From: Dataset
View source
Applies a transformation function to this dataset.
apply enables chaining of custom Dataset transformations, which are represented as functions that take one Dataset argument and return a transformed Dataset.
View source
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.
This method requires that you are running in eager mode and the dataset's element_spec contains only TensorSpec components.
as_numpy_iterator() will preserve the nested structure of dataset elements.
View source
Combines consecutive elements of this dataset into batches.
The components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
View source
Caches the elements in this dataset.
The first time the dataset is iterated over, its elements will be cached either in the specified file or in memory. Subsequent iterations will use the cached data.
When caching to a file, the cached data will persist across runs. Even the first iteration through the data will read from the cache file. Changing the input pipeline before the call to .cache() will have no effect until the cache file is removed or the filename is changed.
View source
Returns the cardinality of the dataset, if known.
cardinality may return tf.data.INFINITE_CARDINALITY if the dataset contains an infinite number of elements or tf.data.UNKNOWN_CARDINALITY if the analysis fails to determine the number of elements in the dataset (e.g. when the dataset source is a file).
View source
Creates a Dataset by concatenating the given dataset with this dataset.
View source
Enumerates the elements of this dataset.
It is similar to python's enumerate.
View source
Filters this dataset according to predicate.
View source
Maps map_func across this dataset and flattens the result.
Use flat_map if you want to make sure that the order of your dataset stays the same. For example, to flatten a dataset of batches into a dataset of their elements:
tf.data.Dataset.interleave() is a generalization of flat_map, since flat_map produces the same output as tf.data.Dataset.interleave(cycle_length=1)
View source
Creates a Dataset whose elements are generated by generator. (deprecated arguments)
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, whichiver 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:
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 with the types defined by output_types and with the shapes which are either unknown or defined by output_shapes.
View source
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.
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.
View source
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.
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.
View source
Maps map_func across this dataset, and interleaves the results.
For example, you can use Dataset.interleave() to process many input files concurrently:
The cycle_length and block_length arguments control the order in which elements are produced. cycle_length controls the number of input elements that are processed concurrently. If you set cycle_length to 1, this transformation will handle one input element at a time, and will produce identical results to tf.data.Dataset.flat_map. In general, this transformation will apply map_func to cycle_length input elements, open iterators on the returned Dataset objects, and cycle through them producing block_length consecutive elements from each iterator, and consuming the next input element each time it reaches the end of an iterator.
Performance can often be improved by setting num_parallel_calls so that interleave will use multiple threads to fetch elements. If determinism isn't required, it can also improve performance to set deterministic=False.
View source
A dataset of all files matching one or more glob patterns.
The file_pattern argument should be a small number of glob patterns. If your filenames have already been globbed, use Dataset.from_tensor_slices(filenames) instead, as re-globbing every filename with list_files may result in poor performance with remote storage systems.
If we had the following files on our filesystem:
If we pass "/path/to/dir/*.py" as the directory, the dataset would produce:
View source
Maps map_func across the elements of this dataset.
This transformation applies map_func to each element of this dataset, and returns a new dataset containing the transformed elements, in the same order as they appeared in the input. map_func can be used to change both the values and the structure of a dataset's elements. For example, adding 1 to each element, or projecting a subset of element components.
The input signature of map_func is determined by the structure of each element in this dataset.
The value or values returned by map_func determine the structure of each element in the returned dataset.
map_func can accept as arguments and return any type of dataset element.
Note that irrespective of the context in which map_func is defined (eager vs. graph), tf.data traces the function and executes it as a graph. To use Python code inside of the function you have a few options:
1) Rely on AutoGraph to convert Python code into an equivalent graph computation. The downside of this approach is that AutoGraph can convert some but not all Python code.
2) Use tf.py_function, which allows you to write arbitrary Python code but will generally result in worse performance than 1). For example:
3) Use tf.numpy_function, which also allows you to write arbitrary Python code. Note that tf.py_function accepts tf.Tensor whereas tf.numpy_function accepts numpy arrays and returns only numpy arrays. For example:
Note that the use of tf.numpy_function and tf.py_function in general precludes the possibility of executing user-defined transformations in parallel (because of Python GIL).
Performance can often be improved by setting num_parallel_calls so that map will use multiple threads to process elements. If deterministic order isn't required, it can also improve performance to set deterministic=False.
View source
Returns the options for this dataset and its inputs.
View source
Combines consecutive elements of this dataset into padded batches.
This transformation combines multiple consecutive elements of the input dataset into a single element.
Like tf.data.Dataset.batch, the components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
Unlike tf.data.Dataset.batch, the input elements to be batched may have different shapes, and this transformation will pad each component to the respective shape in padded_shapes. The padded_shapes argument determines the resulting shape for each dimension of each component in an output element:
See also tf.data.experimental.dense_to_sparse_batch, which combines elements that may have different shapes into a tf.sparse.SparseTensor.
View source
Creates a Dataset that prefetches elements from this dataset.
Most dataset input pipelines should end with a call to prefetch. This allows later elements to be prepared while the current element is being processed. This often improves latency and throughput, at the cost of using additional memory to store prefetched elements.
View source
Creates a Dataset of a step-separated range of values.
View source
Reduces the input dataset to a single element.
The transformation calls reduce_func successively on every element of the input dataset until the dataset is exhausted, aggregating information in its internal state. The initial_state argument is used for the initial state and the final state is returned as the result.
View source
Repeats this dataset so each original value is seen count times.
View source
Creates a Dataset that includes only 1/num_shards of this dataset.
shard is deterministic. The Dataset produced by A.shard(n, i) will contain all elements of A whose index mod n = i.
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 shard elements as follows:
View source
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.
reshuffle_each_iteration controls whether the shuffle order should be different for each epoch. In TF 1.X, the idiomatic way to create epochs was through the repeat transformation:
In TF 2.0, tf.data.Dataset objects are Python iterables which makes it possible to also create epochs through Python iteration:
View source
Creates a Dataset that skips count elements from this dataset.
View source
Creates a Dataset with at most count elements from this dataset.
View source
Splits elements of a dataset into multiple elements.
For example, if elements of the dataset are shaped [B, a0, a1, ...], where B may vary for each input element, then for each element in the dataset, the unbatched dataset will contain B consecutive elements of shape [a0, a1, ...].
View source
Combines (nests of) input elements into a dataset of (nests of) windows.
A "window" is a finite dataset of flat elements of size size (or possibly fewer if there are not enough input elements to fill the window and drop_remainder evaluates to False).
The shift argument determines the number of input elements by which the window moves on each iteration. If windows and elements are both numbered starting at 0, the first element in window k will be element k * shift of the input dataset. In particular, the first element of the first window will always be the first element of the input dataset.
The stride argument determines the stride of the input elements, and the shift argument determines the shift of the window.
Note that when the window transformation is applied to a dataset of nested elements, it produces a dataset of nested windows.
View source
Returns a new tf.data.Dataset with the given options set.
The options are "global" in the sense they apply to the entire dataset. If options are set multiple times, they are merged as long as different options do not use different non-default values.
View source
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.
View source
View source
Creates an iterator for elements of this dataset.
The returned iterator implements the Python Iterator protocol.
View source
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.
View source


A Dataset comprising records from one or more TFRecord files.
Inherits From: Dataset
View source
Applies a transformation function to this dataset.
apply enables chaining of custom Dataset transformations, which are represented as functions that take one Dataset argument and return a transformed Dataset.
View source
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.
This method requires that you are running in eager mode and the dataset's element_spec contains only TensorSpec components.
as_numpy_iterator() will preserve the nested structure of dataset elements.
View source
Combines consecutive elements of this dataset into batches.
The components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
View source
Caches the elements in this dataset.
The first time the dataset is iterated over, its elements will be cached either in the specified file or in memory. Subsequent iterations will use the cached data.
When caching to a file, the cached data will persist across runs. Even the first iteration through the data will read from the cache file. Changing the input pipeline before the call to .cache() will have no effect until the cache file is removed or the filename is changed.
View source
Returns the cardinality of the dataset, if known.
cardinality may return tf.data.INFINITE_CARDINALITY if the dataset contains an infinite number of elements or tf.data.UNKNOWN_CARDINALITY if the analysis fails to determine the number of elements in the dataset (e.g. when the dataset source is a file).
View source
Creates a Dataset by concatenating the given dataset with this dataset.
View source
Enumerates the elements of this dataset.
It is similar to python's enumerate.
View source
Filters this dataset according to predicate.
View source
Maps map_func across this dataset and flattens the result.
Use flat_map if you want to make sure that the order of your dataset stays the same. For example, to flatten a dataset of batches into a dataset of their elements:
tf.data.Dataset.interleave() is a generalization of flat_map, since flat_map produces the same output as tf.data.Dataset.interleave(cycle_length=1)
View source
Creates a Dataset whose elements are generated by generator. (deprecated arguments)
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, whichiver 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:
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 with the types defined by output_types and with the shapes which are either unknown or defined by output_shapes.
View source
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.
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.
View source
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.
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.
View source
Maps map_func across this dataset, and interleaves the results.
For example, you can use Dataset.interleave() to process many input files concurrently:
The cycle_length and block_length arguments control the order in which elements are produced. cycle_length controls the number of input elements that are processed concurrently. If you set cycle_length to 1, this transformation will handle one input element at a time, and will produce identical results to tf.data.Dataset.flat_map. In general, this transformation will apply map_func to cycle_length input elements, open iterators on the returned Dataset objects, and cycle through them producing block_length consecutive elements from each iterator, and consuming the next input element each time it reaches the end of an iterator.
Performance can often be improved by setting num_parallel_calls so that interleave will use multiple threads to fetch elements. If determinism isn't required, it can also improve performance to set deterministic=False.
View source
A dataset of all files matching one or more glob patterns.
The file_pattern argument should be a small number of glob patterns. If your filenames have already been globbed, use Dataset.from_tensor_slices(filenames) instead, as re-globbing every filename with list_files may result in poor performance with remote storage systems.
If we had the following files on our filesystem:
If we pass "/path/to/dir/*.py" as the directory, the dataset would produce:
View source
Maps map_func across the elements of this dataset.
This transformation applies map_func to each element of this dataset, and returns a new dataset containing the transformed elements, in the same order as they appeared in the input. map_func can be used to change both the values and the structure of a dataset's elements. For example, adding 1 to each element, or projecting a subset of element components.
The input signature of map_func is determined by the structure of each element in this dataset.
The value or values returned by map_func determine the structure of each element in the returned dataset.
map_func can accept as arguments and return any type of dataset element.
Note that irrespective of the context in which map_func is defined (eager vs. graph), tf.data traces the function and executes it as a graph. To use Python code inside of the function you have a few options:
1) Rely on AutoGraph to convert Python code into an equivalent graph computation. The downside of this approach is that AutoGraph can convert some but not all Python code.
2) Use tf.py_function, which allows you to write arbitrary Python code but will generally result in worse performance than 1). For example:
3) Use tf.numpy_function, which also allows you to write arbitrary Python code. Note that tf.py_function accepts tf.Tensor whereas tf.numpy_function accepts numpy arrays and returns only numpy arrays. For example:
Note that the use of tf.numpy_function and tf.py_function in general precludes the possibility of executing user-defined transformations in parallel (because of Python GIL).
Performance can often be improved by setting num_parallel_calls so that map will use multiple threads to process elements. If deterministic order isn't required, it can also improve performance to set deterministic=False.
View source
Returns the options for this dataset and its inputs.
View source
Combines consecutive elements of this dataset into padded batches.
This transformation combines multiple consecutive elements of the input dataset into a single element.
Like tf.data.Dataset.batch, the components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
Unlike tf.data.Dataset.batch, the input elements to be batched may have different shapes, and this transformation will pad each component to the respective shape in padded_shapes. The padded_shapes argument determines the resulting shape for each dimension of each component in an output element:
See also tf.data.experimental.dense_to_sparse_batch, which combines elements that may have different shapes into a tf.sparse.SparseTensor.
View source
Creates a Dataset that prefetches elements from this dataset.
Most dataset input pipelines should end with a call to prefetch. This allows later elements to be prepared while the current element is being processed. This often improves latency and throughput, at the cost of using additional memory to store prefetched elements.
View source
Creates a Dataset of a step-separated range of values.
View source
Reduces the input dataset to a single element.
The transformation calls reduce_func successively on every element of the input dataset until the dataset is exhausted, aggregating information in its internal state. The initial_state argument is used for the initial state and the final state is returned as the result.
View source
Repeats this dataset so each original value is seen count times.
View source
Creates a Dataset that includes only 1/num_shards of this dataset.
shard is deterministic. The Dataset produced by A.shard(n, i) will contain all elements of A whose index mod n = i.
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 shard elements as follows:
View source
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.
reshuffle_each_iteration controls whether the shuffle order should be different for each epoch. In TF 1.X, the idiomatic way to create epochs was through the repeat transformation:
In TF 2.0, tf.data.Dataset objects are Python iterables which makes it possible to also create epochs through Python iteration:
View source
Creates a Dataset that skips count elements from this dataset.
View source
Creates a Dataset with at most count elements from this dataset.
View source
Splits elements of a dataset into multiple elements.
For example, if elements of the dataset are shaped [B, a0, a1, ...], where B may vary for each input element, then for each element in the dataset, the unbatched dataset will contain B consecutive elements of shape [a0, a1, ...].
View source
Combines (nests of) input elements into a dataset of (nests of) windows.
A "window" is a finite dataset of flat elements of size size (or possibly fewer if there are not enough input elements to fill the window and drop_remainder evaluates to False).
The shift argument determines the number of input elements by which the window moves on each iteration. If windows and elements are both numbered starting at 0, the first element in window k will be element k * shift of the input dataset. In particular, the first element of the first window will always be the first element of the input dataset.
The stride argument determines the stride of the input elements, and the shift argument determines the shift of the window.
Note that when the window transformation is applied to a dataset of nested elements, it produces a dataset of nested windows.
View source
Returns a new tf.data.Dataset with the given options set.
The options are "global" in the sense they apply to the entire dataset. If options are set multiple times, they are merged as long as different options do not use different non-default values.
View source
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.
View source
View source
Creates an iterator for elements of this dataset.
The returned iterator implements the Python Iterator protocol.
View source
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.
View source


Public API for tf.debugging namespace.
experimental module: Public API for tf.debugging.experimental namespace.
Assert(...): Asserts that the given condition is true.
assert_all_finite(...): Assert that the tensor does not contain any NaN's or Inf's.
assert_equal(...): Assert the condition x == y holds element-wise.
assert_greater(...): Assert the condition x > y holds element-wise.
assert_greater_equal(...): Assert the condition x >= y holds element-wise.
assert_integer(...): Assert that x is of integer dtype.
assert_less(...): Assert the condition x < y holds element-wise.
assert_less_equal(...): Assert the condition x <= y holds element-wise.
assert_near(...): Assert the condition x and y are close element-wise.
assert_negative(...): Assert the condition x < 0 holds element-wise.
assert_non_negative(...): Assert the condition x >= 0 holds element-wise.
assert_non_positive(...): Assert the condition x <= 0 holds element-wise.
assert_none_equal(...): Assert the condition x != y holds for all elements.
assert_positive(...): Assert the condition x > 0 holds element-wise.
assert_proper_iterable(...): Static assert that values is a "proper" iterable.
assert_rank(...): Assert that x has rank equal to rank.
assert_rank_at_least(...): Assert that x has rank of at least rank.
assert_rank_in(...): Assert that x has a rank in ranks.
assert_same_float_dtype(...): Validate and return float type based on tensors and dtype.
assert_scalar(...): Asserts that the given tensor is a scalar.
assert_shapes(...): Assert tensor shapes and dimension size relationships between tensors.
assert_type(...): Asserts that the given Tensor is of the specified type.
check_numerics(...): Checks a tensor for NaN and Inf values.
disable_check_numerics(...): Disable the eager/graph unified numerics checking mechanism.
enable_check_numerics(...): Enable tensor numerics checking in an eager/graph unified fashion.
get_log_device_placement(...): Get if device placements are logged.
is_numeric_tensor(...): Returns True if the elements of tensor are numbers.
set_log_device_placement(...): Set if device placements should be logged.


Asserts that the given condition is true.
Main aliases
tf.Assert
See Migration guide for more details.
tf.compat.v1.Assert, tf.compat.v1.debugging.Assert
If condition evaluates to false, print the list of tensors in data. summarize determines how many entries of the tensors to print.
When in TF V1 mode (that is, outside tf.function) Assert needs a control dependency on the output to ensure the assertion executes:
returns None


Assert that the tensor does not contain any NaN's or Inf's.


Assert the condition x == y holds element-wise.
Main aliases
tf.assert_equal
This Op checks that x[i] == y[i] holds for every pair of (possibly broadcast) elements of x and y. If both x and y are empty, this is trivially satisfied.
If x and y are not equal, message, as well as the first summarize entries of x and y are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x > y holds element-wise.
Main aliases
tf.assert_greater
This Op checks that x[i] > y[i] holds for every pair of (possibly broadcast) elements of x and y. If both x and y are empty, this is trivially satisfied.
If x is not greater than y element-wise, message, as well as the first summarize entries of x and y are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x >= y holds element-wise.
This Op checks that x[i] >= y[i] holds for every pair of (possibly broadcast) elements of x and y. If both x and y are empty, this is trivially satisfied.
If x is not greater or equal to y element-wise, message, as well as the first summarize entries of x and y are printed, and InvalidArgumentError is raised.
returns None


Assert that x is of integer dtype.
If x has a non-integer type, message, as well as the dtype of x are printed, and InvalidArgumentError is raised.
This can always be checked statically, so this method returns nothing.


Assert the condition x < y holds element-wise.
Main aliases
tf.assert_less
This Op checks that x[i] < y[i] holds for every pair of (possibly broadcast) elements of x and y. If both x and y are empty, this is trivially satisfied.
If x is not less than y element-wise, message, as well as the first summarize entries of x and y are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x <= y holds element-wise.
This Op checks that x[i] <= y[i] holds for every pair of (possibly broadcast) elements of x and y. If both x and y are empty, this is trivially satisfied.
If x is not less or equal than y element-wise, message, as well as the first summarize entries of x and y are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x and y are close element-wise.
This Op checks that x[i] - y[i] < atol + rtol * tf.abs(y[i]) holds for every pair of (possibly broadcast) elements of x and y. If both x and y are empty, this is trivially satisfied.
If any elements of x and y are not close, message, as well as the first summarize entries of x and y are printed, and InvalidArgumentError is raised.
The default atol and rtol is 10 * eps, where eps is the smallest representable positive number such that 1 + eps != 1. This is about 1.2e-6 in 32bit, 2.22e-15 in 64bit, and 0.00977 in 16bit. See numpy.finfo.
returns None
Similar to numpy.testing.assert_allclose, except tolerance depends on data type. This is due to the fact that TensorFlow is often used with 32bit, 64bit, and even 16bit data.


Assert the condition x < 0 holds element-wise.
This Op checks that x[i] < 0 holds for every element of x. If x is empty, this is trivially satisfied.
If x is not negative everywhere, message, as well as the first summarize entries of x are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x != y holds for all elements.
This Op checks that x[i] != y[i] holds for every pair of (possibly broadcast) elements of x and y. If both x and y are empty, this is trivially satisfied.
If any elements of x and y are equal, message, as well as the first summarize entries of x and y are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x >= 0 holds element-wise.
This Op checks that x[i] >= 0 holds for every element of x. If x is empty, this is trivially satisfied.
If x is not >= 0 everywhere, message, as well as the first summarize entries of x are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x <= 0 holds element-wise.
This Op checks that x[i] <= 0 holds for every element of x. If x is empty, this is trivially satisfied.
If x is not <= 0 everywhere, message, as well as the first summarize entries of x are printed, and InvalidArgumentError is raised.
returns None


Assert the condition x > 0 holds element-wise.
This Op checks that x[i] > 0 holds for every element of x. If x is empty, this is trivially satisfied.
If x is not positive everywhere, message, as well as the first summarize entries of x are printed, and InvalidArgumentError is raised.
returns None


Static assert that values is a "proper" iterable.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.assert_proper_iterable, tf.compat.v1.debugging.assert_proper_iterable
Ops that expect iterables of Tensor can call this to validate input. Useful since Tensor, ndarray, byte/text type are all iterables themselves.


Assert that x has rank equal to rank.
Main aliases
tf.assert_rank
This Op checks that the rank of x is equal to rank.
If x has a different rank, message, as well as the shape of x are printed, and InvalidArgumentError is raised.
returns None


Assert that x has rank of at least rank.
This Op checks that the rank of x is greater or equal to rank.
If x has a rank lower than rank, message, as well as the shape of x are printed, and InvalidArgumentError is raised.
returns None


Assert that x has a rank in ranks.
This Op checks that the rank of x is in ranks.
If x has a different rank, message, as well as the shape of x are printed, and InvalidArgumentError is raised.
returns None


Validate and return float type based on tensors and dtype.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.assert_same_float_dtype, tf.compat.v1.debugging.assert_same_float_dtype
For ops such as matrix multiplication, inputs and weights must be of the same float type. This function validates that all tensors are the same type, validates that type is dtype (if supplied), and returns the type. Type must be a floating point type. If neither tensors nor dtype is supplied, the function will return dtypes.float32.


Asserts that the given tensor is a scalar.
This function raises ValueError unless it can be certain that the given tensor is a scalar. ValueError is also raised if the shape of tensor is unknown.
This is always checked statically, so this method returns nothing.


Assert tensor shapes and dimension size relationships between tensors.
This Op checks that a collection of tensors shape relationships satisfies given constraints.
If x, y, param or scalar does not have a shape that satisfies all specified constraints, message, as well as the first summarize entries of the first encountered violating tensor are printed, and InvalidArgumentError is raised.
Size entries in the specified shapes are checked against other entries by their hash, except:
If the first entry of a shape is ... (type Ellipsis) or '*' that indicates a variable number of outer dimensions of unspecified size, i.e. the constraint applies to the inner-most dimensions only.
Scalar tensors and specified shapes of length zero (excluding the 'inner-most' prefix) are both treated as having a single dimension of size one.


Asserts that the given Tensor is of the specified type.
This can always be checked statically, so this method returns nothing.


Checks a tensor for NaN and Inf values.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.check_numerics, tf.compat.v1.debugging.check_numerics
When run, reports an InvalidArgument error if tensor has any values that are not a number (NaN) or infinity (Inf). Otherwise, passes tensor as-is.


Disable the eager/graph unified numerics checking mechanism.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.debugging.disable_check_numerics
This method can be used after a call to tf.debugging.enable_check_numerics() to disable the numerics-checking mechanism that catches infinity and NaN values output by ops executed eagerly or in tf.function-compiled graphs.
This method is idempotent. Calling it multiple times has the same effect as calling it once.
This method takes effect only on the thread in which it is called.


Enable tensor numerics checking in an eager/graph unified fashion.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.debugging.enable_check_numerics
The numerics checking mechanism will cause any TensorFlow eager execution or graph execution to error out as soon as an op's output tensor contains infinity or NaN.
This method is idempotent. Calling it multiple times has the same effect as calling it once.
This method takes effect only on the thread in which it is called.
When a op's float-type output tensor contains any Infinity or NaN, an tf.errors.InvalidArgumentError will be thrown, with an error message that reveals the following information:
Once enabled, the check-numerics mechanism can be disabled by using tf.debugging.disable_check_numerics().
Catching infinity during the execution of a tf.function graph:
Catching NaN during eager execution:


Public API for tf.debugging.experimental namespace.
disable_dump_debug_info(...): Disable the currently-enabled debugging dumping.
enable_dump_debug_info(...): Enable dumping debugging information from a TensorFlow program.


Disable the currently-enabled debugging dumping.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.debugging.experimental.disable_dump_debug_info
If the enable_dump_debug_info() method under the same Python namespace has been invoked before, calling this method disables it. If no call to enable_dump_debug_info() has been made, calling this method is a no-op. Calling this method more than once is idempotent.


Enable dumping debugging information from a TensorFlow program.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.debugging.experimental.enable_dump_debug_info
The debugging information is dumped to a directory on the file system specified as dump_root.
The dumped debugging information can be ingested by debugger UIs.
The files in the dump directory contain the following information:
Once enabled, the dumping can be disabled with the corresponding disable_dump_debug_info() method under the same Python namespace. Calling this method more than once with the same dump_root is idempotent. Calling this method more than once with different tensor_debug_modes leads to a ValueError. Calling this method more than once with different circular_buffer_sizes leads to a ValueError. Calling this method with a different dump_root abolishes the previously-enabled dump_root.


Get if device placements are logged.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.debugging.get_log_device_placement


Returns True if the elements of tensor are numbers.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.debugging.is_numeric_tensor, tf.compat.v1.is_numeric_tensor
Specifically, returns True if the dtype of tensor is one of the following:
Returns False if tensor is of a non-numeric type or if tensor is not a tf.Tensor object.


Set if device placements should be logged.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.debugging.set_log_device_placement


Specifies the device for ops created/executed in this context.
This function specifies the device to be used for ops created/executed in a particular context. Nested contexts will inherit and also create/execute their ops on the specified device. If a specific device is not required, consider not using this function so that a device can be automatically assigned. In general the use of this function is optional. device_name can be fully specified, as in "/job:worker/task:1/device:cpu:0", or partially specified, containing only a subset of the "/"-separated fields. Any fields which are specified will override device annotations from outer scopes.


Represents a (possibly partial) specification for a TensorFlow device.
DeviceSpecs are used throughout TensorFlow to describe where state is stored and computations occur. Using DeviceSpec allows you to parse device spec strings to verify their validity, merge them or compose them programmatically.
With eager execution disabled (by default in TensorFlow 1.x and by calling disable_eager_execution() in TensorFlow 2.x), the following syntax can be used:
A DeviceSpec consists of 5 components -- each of which is optionally specified:
View source
Construct a DeviceSpec from a string.
View source
Returns a new DeviceSpec which incorporates dev.
When combining specs, dev will take precedence over the current spec. So for instance:
is equivalent to:
View source
Parse a DeviceSpec name into its components.
2.x behavior change: In TensorFlow 1.x, this function mutates its own state and returns itself. In 2.x, DeviceSpecs are immutable, and this function will return a DeviceSpec which contains the spec.
Recommended:
Will work in 1.x and 2.x (though deprecated in 2.x):
Will NOT work in 2.x:
In general, DeviceSpec.from_string should completely replace DeviceSpec.parse_from_string, and DeviceSpec.replace should completely replace setting attributes directly.
View source
Convenience method for making a new DeviceSpec by overriding fields.
View source
Return a string representation of this DeviceSpec.
View source
Checks if the other DeviceSpec is same as the current instance, eg have
same value for all the internal fields.


Library for running a computation across multiple devices.
The intent of this library is that you can write an algorithm in a stylized way and it will be usable with a variety of different tf.distribute.Strategy implementations. Each descendant will implement a different strategy for distributing the algorithm across multiple devices/machines. Furthermore, these changes can be hidden inside the specific layers and other library classes that need special treatment to run in a distributed setting, so that most users' model definition code can run unchanged. The tf.distribute.Strategy API works the same way with eager and graph execution.
Guides
Tutorials
Distributed Training Tutorials
The tutorials cover how to use tf.distribute.Strategy to do distributed training with native Keras APIs, custom training loops, and Esitmator APIs. They also cover how to save/load model when using tf.distribute.Strategy.
Glossary
Parameter servers: These are machines that hold a single copy of parameters/variables, used by some strategies (right now just tf.distribute.experimental.ParameterServerStrategy). All replicas that want to operate on a variable retrieve it at the beginning of a step and send an update to be applied at the end of the step. These can in priniciple support either sync or async training, but right now we only have support for async training with parameter servers. Compare to tf.distribute.experimental.CentralStorageStrategy, which puts all variables on a single device on the same machine (and does sync training), and tf.distribute.MirroredStrategy, which mirrors variables to multiple devices (see below).
Replica context vs. Cross-replica context vs Update context
A replica context applies when you execute the computation function that was called with strategy.run. Conceptually, you're in replica context when executing the computation function that is being replicated.
An update context is entered in a tf.distribute.StrategyExtended.update call.
An cross-replica context is entered when you enter a strategy.scope. This is useful for calling tf.distribute.Strategy methods which operate across the replicas (like reduce_to()). By default you start in a replica context (the "default single replica context") and then some methods can switch you back and forth.
Distributed value: Distributed value is represented by the base class tf.distribute.DistributedValues. tf.distribute.DistributedValues is useful to represent values on multiple devices, and it contains a map from replica id to values. Two representative kinds of tf.distribute.DistributedValues are "PerReplica" and "Mirrored" values.
"PerReplica" values exist on the worker devices, with a different value for each replica. They are produced by iterating through a distributed dataset returned by tf.distribute.Strategy.experimental_distribute_dataset and tf.distribute.Strategy.distribute_datasets_from_function. They are also the typical result returned by tf.distribute.Strategy.run.
"Mirrored" values are like "PerReplica" values, except we know that the value on all replicas are the same. We can safely read a "Mirrored" value in a cross-replica context by using the value on any replica.
Unwrapping and merging: Consider calling a function fn on multiple replicas, like strategy.run(fn, args=[w]) with an argument w that is a tf.distribute.DistributedValues. This means w will have a map taking replica id 0 to w0, replica id 1 to w1, etc. strategy.run() unwraps w before calling fn, so it calls fn(w0) on device d0, fn(w1) on device d1, etc. It then merges the return values from fn(), which leads to one common object if the returned values are the same object from every replica, or a DistributedValues object otherwise.
Reductions and all-reduce: A reduction is a method of aggregating multiple values into one value, like "sum" or "mean". If a strategy is doing sync training, we will perform a reduction on the gradients to a parameter from all replicas before applying the update. All-reduce is an algorithm for performing a reduction on values from multiple devices and making the result available on all of those devices.
Mirrored variables: These are variables that are created on multiple devices, where we keep the variables in sync by applying the same updates to every copy. Mirrored variables are created with tf.Variable(...synchronization=tf.VariableSynchronization.ON_WRITE...). Normally they are only used in synchronous training.
SyncOnRead variables
SyncOnRead variables are created by tf.Variable(...synchronization=tf.VariableSynchronization.ON_READ...), and they are created on multiple devices. In replica context, each component variable on the local replica can perform reads and writes without synchronization with each other. When the SyncOnRead variable is read in cross-replica context, the values from component variables are aggregated and returned.
SyncOnRead variables bring a lot of custom configuration difficulty to the underlying logic, so we do not encourage users to instantiate and use SyncOnRead variable on their own. We have mainly used SyncOnRead variables for use cases such as batch norm and metrics. For performance reasons, we often don't need to keep these statistics in sync every step and they can be accumulated on each replica independently. The only time we want to sync them is reporting or checkpointing, which typically happens in cross-replica context. SyncOnRead variables are also often used by advanced users who want to control when variable values are aggregated. For example, users sometimes want to maintain gradients independently on each replica for a couple of steps without aggregation.
Distribute-aware layers
Layers are generally called in a replica context, except when defining a Keras functional model. tf.distribute.in_cross_replica_context will let you determine which case you are in. If in a replica context, the tf.distribute.get_replica_context function will return the default replica context outside a strategy scope, None within a strategy scope, and a tf.distribute.ReplicaContext object inside a strategy scope and within a tf.distribute.Strategy.run function. The ReplicaContext object has an all_reduce method for aggregating across all replicas.
Note that we provide a default version of tf.distribute.Strategy that is used when no other strategy is in scope, that provides the same API with reasonable default behavior.
cluster_resolver module: Library imports for ClusterResolvers.
experimental module: Public API for tf.distribute.experimental namespace.
class CrossDeviceOps: Base class for cross-device reduction and broadcasting algorithms.
class DistributedDataset: Represents a dataset distributed among devices and machines.
class DistributedIterator: An iterator over tf.distribute.DistributedDataset.
class DistributedValues: Base class for representing distributed values.
class HierarchicalCopyAllReduce: Hierarchical copy all-reduce implementation of CrossDeviceOps.
class InputContext: A class wrapping information needed by an input function.
class InputOptions: Run options for experimental_distribute_dataset(s_from_function).
class InputReplicationMode: Replication mode for input function.
class MirroredStrategy: Synchronous training across multiple replicas on one machine.
class MultiWorkerMirroredStrategy: A distribution strategy for synchronous training on multiple workers.
class NcclAllReduce: NCCL all-reduce implementation of CrossDeviceOps.
class OneDeviceStrategy: A distribution strategy for running on a single device.
class ReduceOp: Indicates how a set of values should be reduced.
class ReductionToOneDevice: A CrossDeviceOps implementation that copies values to one device to reduce.
class ReplicaContext: A class with a collection of APIs that can be called in a replica context.
class RunOptions: Run options for strategy.run.
class Server: An in-process TensorFlow server, for use in distributed training.
class Strategy: A state & compute distribution policy on a list of devices.
class StrategyExtended: Additional APIs for algorithms that need to be distribution-aware.
class TPUStrategy: Synchronous training on TPUs and TPU Pods.
experimental_set_strategy(...): Set a tf.distribute.Strategy as current without with strategy.scope().
get_replica_context(...): Returns the current tf.distribute.ReplicaContext or None.
get_strategy(...): Returns the current tf.distribute.Strategy object.
has_strategy(...): Return if there is a current non-default tf.distribute.Strategy.
in_cross_replica_context(...): Returns True if in a cross-replica context.


Library imports for ClusterResolvers.
This library contains all implementations of ClusterResolvers. ClusterResolvers are a way of specifying cluster information for distributed execution. Built on top of existing ClusterSpec framework, ClusterResolvers are a way for TensorFlow to communicate with various cluster management systems (e.g. GCE, AWS, etc...).
class ClusterResolver: Abstract class for all implementations of ClusterResolvers.
class GCEClusterResolver: ClusterResolver for Google Compute Engine.
class KubernetesClusterResolver: ClusterResolver for Kubernetes.
class SimpleClusterResolver: Simple implementation of ClusterResolver that accepts all attributes.
class SlurmClusterResolver: ClusterResolver for system with Slurm workload manager.
class TFConfigClusterResolver: Implementation of a ClusterResolver which reads the TF_CONFIG EnvVar.
class TPUClusterResolver: Cluster Resolver for Google Cloud TPUs.
class UnionResolver: Performs a union on underlying ClusterResolvers.


Abstract class for all implementations of ClusterResolvers.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.ClusterResolver
This defines the skeleton for all implementations of ClusterResolvers. ClusterResolvers are a way for TensorFlow to communicate with various cluster management systems (e.g. GCE, AWS, etc...) and gives TensorFlow necessary information to set up distributed training.
By letting TensorFlow communicate with these systems, we will be able to automatically discover and resolve IP addresses for various TensorFlow workers. This will eventually allow us to automatically recover from underlying machine failures and scale TensorFlow worker clusters up and down.
Note to Implementors of tf.distribute.cluster_resolver.ClusterResolver subclass: In addition to these abstract methods, when task_type, task_id, and rpc_layer attributes are applicable, you should also implement them either as properties with getters or setters, or directly set the attributes self._task_type, self._task_id, or self._rpc_layer so the base class' getters and setters are used. See tf.distribute.clusterresolver.SimpleClusterResolver.init_ for an example.
In general, multi-client tf.distribute strategies such as tf.distribute.experimental.MultiWorkerMirroredStrategy require task_type and task_id properties to be available in the ClusterResolver they are using. On the other hand, these concepts are not applicable in single-client strategies, such as tf.distribute.experimental.TPUStrategy, because the program is only expected to be run on one task, so there should not be a need to have code branches according to task type and task id.
There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).
If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.
Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Retrieve the current state of the cluster and return a tf.train.ClusterSpec.
Implementors of this function must take care in ensuring that the ClusterSpec returned is up-to-date at the time of calling this function. This usually means retrieving the information from the underlying cluster management system every time this function is invoked and reconstructing a cluster_spec, rather than attempting to cache anything.
View source
Retrieves the name or URL of the session master.
Implementors of this function must take care in ensuring that the master returned is up-to-date at the time to calling this function. This usually means retrieving the master every time this function is invoked.
View source
Returns the number of accelerator cores per worker.
This returns the number of accelerator cores (such as GPUs and TPUs) available per worker.
Optionally, we allow callers to specify the task_type, and task_id, for if they want to target a specific TensorFlow task to query the number of accelerators. This is to support heterogenous environments, where the number of accelerators cores per host is different.


ClusterResolver for Google Compute Engine.
Inherits From: ClusterResolver
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.GCEClusterResolver
This is an implementation of cluster resolvers for the Google Compute Engine instance group platform. By specifying a project, zone, and instance group, this will retrieve the IP address of all the instances within the instance group and return a ClusterResolver object suitable for use for distributed TensorFlow.
Usage example with tf.distribute.Strategy:
There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).
If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.
Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Returns a ClusterSpec object based on the latest instance group info.
This returns a ClusterSpec object for use based on information from the specified instance group. We will retrieve the information from the GCE APIs every time this method is called.
View source
Retrieves the name or URL of the session master.
Implementors of this function must take care in ensuring that the master returned is up-to-date at the time to calling this function. This usually means retrieving the master every time this function is invoked.
View source
Returns the number of accelerator cores per worker.
This returns the number of accelerator cores (such as GPUs and TPUs) available per worker.
Optionally, we allow callers to specify the task_type, and task_id, for if they want to target a specific TensorFlow task to query the number of accelerators. This is to support heterogenous environments, where the number of accelerators cores per host is different.


ClusterResolver for Kubernetes.
Inherits From: ClusterResolver
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.KubernetesClusterResolver
This is an implementation of cluster resolvers for Kubernetes. When given the the Kubernetes namespace and label selector for pods, we will retrieve the pod IP addresses of all running pods matching the selector, and return a ClusterSpec based on that information.
Usage example with tf.distribute.Strategy:
There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).
If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.
Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Returns a ClusterSpec object based on the latest info from Kubernetes.
We retrieve the information from the Kubernetes master every time this method is called.
View source
Returns the master address to use when creating a session.
You must have set the task_type and task_id object properties before calling this function, or pass in the task_type and task_id parameters when using this function. If you do both, the function parameters will override the object properties.
View source
Returns the number of accelerator cores per worker.
This returns the number of accelerator cores (such as GPUs and TPUs) available per worker.
Optionally, we allow callers to specify the task_type, and task_id, for if they want to target a specific TensorFlow task to query the number of accelerators. This is to support heterogenous environments, where the number of accelerators cores per host is different.


Simple implementation of ClusterResolver that accepts all attributes.
Inherits From: ClusterResolver
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.SimpleClusterResolver
Please see the base class for documentation of arguments of its constructor.
It is useful if you want to specify some or all attributes.
Usage example with tf.distribute.Strategy:
There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).
If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.
Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Returns the ClusterSpec passed into the constructor.
View source
Returns the master address to use when creating a session.
If a task_type and task_id is given, this will override the master string passed into the initialization function.
View source
Returns the number of accelerator cores per worker.
The SimpleClusterResolver does not do automatic detection of accelerators, and thus all arguments are unused and we simply return the value provided in the constructor.


ClusterResolver for system with Slurm workload manager.
Inherits From: ClusterResolver
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.SlurmClusterResolver
This is an implementation of ClusterResolver for Slurm clusters. This allows the specification of jobs and task counts, number of tasks per node, number of GPUs on each node and number of GPUs for each task. It retrieves system attributes by Slurm environment variables, resolves allocated computing node names, constructs a cluster and returns a ClusterResolver object which can be used for distributed TensorFlow.
There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).
If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.
Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Returns a ClusterSpec object based on the latest instance group info.
This returns a ClusterSpec object for use based on information from the specified initialization parameters and Slurm environment variables. The cluster specification is resolved each time this function is called. The resolver extract hostnames of nodes by scontrol and pack tasks in that order until a node a has number of tasks that is equal to specification. GPUs on nodes are allocated to tasks by specification through setting CUDA_VISIBLE_DEVICES environment variable.
View source
Returns job name and task_id for the process which calls this.
This returns the job name and task index for the process which calls this function according to its rank and cluster specification. The job name and task index are set after a cluster is constructed by cluster_spec otherwise defaults to None.
View source
Returns the master string for connecting to a TensorFlow master.
View source
Returns the number of accelerator cores per worker.
This returns the number of accelerator cores (such as GPUs and TPUs) available per worker.
Optionally, we allow callers to specify the task_type, and task_id, for if they want to target a specific TensorFlow task to query the number of accelerators. This is to support heterogenous environments, where the number of accelerators cores per host is different.


Implementation of a ClusterResolver which reads the TF_CONFIG EnvVar.
Inherits From: ClusterResolver
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.TFConfigClusterResolver
This is an implementation of cluster resolvers when using TF_CONFIG to set information about the cluster. The cluster spec returned will be initialized from the TF_CONFIG environment variable.
An example to set TF_CONFIG is:
However, sometimes the container orchestration framework will set TF_CONFIG for you. In this case, you can just create an instance without passing in any arguments. You can find an example here to let Kuburnetes set TF_CONFIG for you: https://github.com/tensorflow/ecosystem/tree/master/kubernetes. Then you can use it with tf.distribute.Strategy as:
There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).
If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.
Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Returns a ClusterSpec based on the TF_CONFIG environment variable.
View source
Returns the master address to use when creating a TensorFlow session.
View source
Returns the number of accelerator cores per worker.
This returns the number of accelerator cores (such as GPUs and TPUs) available per worker.
Optionally, we allow callers to specify the task_type, and task_id, for if they want to target a specific TensorFlow task to query the number of accelerators. This is to support heterogenous environments, where the number of accelerators cores per host is different.


Cluster Resolver for Google Cloud TPUs.
Inherits From: ClusterResolver
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.TPUClusterResolver
This is an implementation of cluster resolvers for the Google Cloud TPU service.
TPUClusterResolver supports the following distinct environments: Google Compute Engine Google Kubernetes Engine Google internal
It can be passed into tf.distribute.TPUStrategy to support TF2 training on Cloud TPUs.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Returns a ClusterSpec object based on the latest TPU information.
We retrieve the information from the GCE APIs every time this method is called.
View source
Initializes TPU and returns a TPUClusterResolver.
This API will connect to remote TPU cluster and initialize the TPU hardwares. Example usage:
It can be viewed as convenient wrapper of the following code:
View source
View source
View source
Returns the metadata of the TPU system.
Users can call this method to get some facts of the TPU system, like total number of cores, number of TPU workers and the devices. E.g.
View source
Get the Master string to be used for the session.
In the normal case, this returns the grpc path (grpc://1.2.3.4:8470) of first instance in the ClusterSpec returned by the cluster_spec function.
If a non-TPU name is used when constructing a TPUClusterResolver, that will be returned instead (e.g. If the tpus argument's value when constructing this TPUClusterResolver was 'grpc://10.240.1.2:8470', 'grpc://10.240.1.2:8470' will be returned).
View source
Returns the number of TPU cores per worker.
Connects to the master and list all the devices present in the master, and counts them up. Also verifies that the device counts per host in the cluster is the same before returning the number of TPU cores per host.
View source
View source


Performs a union on underlying ClusterResolvers.
Inherits From: ClusterResolver
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.cluster_resolver.UnionResolver
This class performs a union given two or more existing ClusterResolvers. It merges the underlying ClusterResolvers, and returns one unified ClusterSpec when cluster_spec is called. The details of the merge function is documented in the cluster_spec function.
For additional ClusterResolver properties such as task type, task index, rpc layer, environment, etc..., we will return the value from the first ClusterResolver in the union.
An example to combine two cluster resolvers:
An example to override the number of GPUs in a TFConfigClusterResolver instance:
There are two possible return values, "google" (when TensorFlow is running in a Google-internal environment) or an empty string (when TensorFlow is running elsewhere).
If you are implementing a ClusterResolver that works in both the Google environment and the open-source world (for instance, a TPU ClusterResolver or similar), you will have to return the appropriate string depending on the environment, which you will have to detect.
Otherwise, if you are implementing a ClusterResolver that will only work in open-source TensorFlow, you do not need to implement this property.
In TensorFlow distributed environment, each job may have an applicable task id, which is the index of the instance within its task type. This is useful when user needs to run specific code according to task index. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.cluster_resolver.TPUClusterResolver.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class docstring.
In TensorFlow distributed environment, each job may have an applicable task type. Valid task types in TensorFlow include 'chief': a worker that is designated with more responsibility, 'worker': a regular worker for training/evaluation, 'ps': a parameter server, or 'evaluator': an evaluator that evaluates the checkpoints for metrics.
See Multi-worker configuration for more information about 'chief' and 'worker' task type, which are most commonly used.
Having access to such information is useful when user needs to run specific code according to task types. For example,
Returns None if such information is not available or is not applicable in the current distributed environment, such as training with tf.distribute.experimental.TPUStrategy.
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's class doc.
View source
Returns a union of all the ClusterSpecs from the ClusterResolvers.
If all underlying ClusterSpecs expose the set of workers as lists, we will concatenate the lists of workers, starting with the list of workers from the first ClusterResolver passed into the constructor.
If any of the ClusterSpecs expose the set of workers as a dict, we will treat all the sets of workers as dicts (even if they are returned as lists) and will only merge them into a dict if there is no conflicting keys. If there is a conflicting key, we will raise a KeyError.
View source
Returns the master address to use when creating a session.
This usually returns the master from the first ClusterResolver passed in, but you can override this by specifying the task_type and task_id.
View source
Returns the number of accelerator cores per worker.
This returns the number of accelerator cores (such as GPUs and TPUs) available per worker.
Optionally, we allow callers to specify the task_type, and task_id, for if they want to target a specific TensorFlow task to query the number of accelerators. This is to support heterogenous environments, where the number of accelerators cores per host is different.


Base class for cross-device reduction and broadcasting algorithms.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.CrossDeviceOps
The main purpose of this class is to be passed to tf.distribute.MirroredStrategy in order to choose among different cross device communication implementations. Prefer using the methods of tf.distribute.Strategy instead of the ones of this class.
View source
Reduce values to destinations in batches.
See tf.distribute.StrategyExtended.batch_reduce_to. This can only be called in the cross-replica context.
View source
Implementation of batch_reduce.
Overriding this method is useful for subclass implementers.
View source
Broadcast tensor to destinations.
This can only be called in the cross-replica context.
View source
Implementation of broadcast.
View source
Reduce per_replica_value to destinations.
See tf.distribute.StrategyExtended.reduce_to. This can only be called in the cross-replica context.
View source
Implementation of reduce.
Overriding this method is useful for subclass implementers.


Represents a dataset distributed among devices and machines.
A tf.distribute.DistributedDataset could be thought of as a "distributed" dataset. When you use tf.distribute API to scale training to multiple devices or machines, you also need to distribute the input data, which leads to a tf.distribute.DistributedDataset instance, instead of a tf.data.Dataset instance in the non-distributed case. In TF 2.x, tf.distribute.DistributedDataset objects are Python iterables.
There are two APIs to create a tf.distribute.DistributedDataset object: tf.distribute.Strategy.experimental_distribute_dataset(dataset)and tf.distribute.Strategy.distribute_datasets_from_function(dataset_fn). When to use which? When you have a tf.data.Dataset instance, and the regular batch splitting (i.e. re-batch the input tf.data.Dataset instance with a new batch size that is equal to the global batch size divided by the number of replicas in sync) and autosharding (i.e. the tf.data.experimental.AutoShardPolicy options) work for you, use the former API. Otherwise, if you are not using a canonical tf.data.Dataset instance, or you would like to customize the batch splitting or sharding, you can wrap these logic in a dataset_fn and use the latter API. Both API handles prefetch to device for the user. For more details and examples, follow the links to the APIs.
There are two main usages of a DistributedDataset object:
Iterate over it to generate the input for a single device or multiple devices, which is a tf.distribute.DistributedValues instance. To do this, you can:
To achieve a performance improvement, you can also wrap the strategy.run call with a tf.range inside a tf.function. This runs multiple steps in a tf.function. Autograph will convert it to a tf.while_loop on the worker. However, it is less flexible comparing with running a single step inside tf.function. For example, you cannot run things eagerly or arbitrary python code within the steps.
Inspect the tf.TypeSpec of the data generated by DistributedDataset.
tf.distribute.DistributedDataset generates tf.distribute.DistributedValues as input to the devices. If you pass the input to a tf.function and would like to specify the shape and type of each Tensor argument to the function, you can pass a tf.TypeSpec object to the input_signature argument of the tf.function. To get the tf.TypeSpec of the input, you can use the element_spec property of the tf.distribute.DistributedDataset or tf.distribute.DistributedIterator object.
For example:
Visit the tutorial on distributed input for more examples and caveats.
View source
Creates an iterator for the tf.distribute.DistributedDataset.
The returned iterator implements the Python Iterator protocol.


An iterator over tf.distribute.DistributedDataset.
tf.distribute.DistributedIterator is the primary mechanism for enumerating elements of a tf.distribute.DistributedDataset. It supports the Python Iterator protocol, which means it can be iterated over using a for-loop or by fetching individual elements explicitly via get_next().
You can create a tf.distribute.DistributedIterator by calling iter on a tf.distribute.DistributedDataset or creating a python loop over a tf.distribute.DistributedDataset.
Visit the tutorial on distributed input for more examples and caveats.
View source
Returns the next input from the iterator for all replicas.
View source
Returns a tf.experimental.Optional that contains the next value for all replicas.
If the tf.distribute.DistributedIterator has reached the end of the sequence, the returned tf.experimental.Optional will have no value.


Base class for representing distributed values.
A subclass instance of tf.distribute.DistributedValues is created when creating variables within a distribution strategy, iterating a tf.distribute.DistributedDataset or through tf.distribute.Strategy.run. This base class should never be instantiated directly. tf.distribute.DistributedValues contains a value per replica. Depending on the subclass, the values could either be synced on update, synced on demand, or never synced.
tf.distribute.DistributedValues can be reduced to obtain single value across replicas, as input into tf.distribute.Strategy.run or the per-replica values inspected using tf.distribute.Strategy.experimental_local_results.


Public API for tf.distribute.experimental namespace.
coordinator module: Public API for tf.distribute.experimental.coordinator namespace.
partitioners module: Public API for tf.distribute.experimental.partitioners namespace.
class CentralStorageStrategy: A one-machine strategy that puts all variables on a single device.
class CollectiveCommunication: Cross device communication implementation.
class CollectiveHints: Hints for collective operations like AllReduce.
class CommunicationImplementation: Cross device communication implementation.
class CommunicationOptions: Options for cross device communications like All-reduce.
class MultiWorkerMirroredStrategy: A distribution strategy for synchronous training on multiple workers.
class ParameterServerStrategy: An multi-worker tf.distribute strategy with parameter servers.
class TPUStrategy: Synchronous training on TPUs and TPU Pods.
class ValueContext: A class wrapping information needed by a distribute function.


A one-machine strategy that puts all variables on a single device.
Inherits From: Strategy
Variables are assigned to local CPU or the only GPU. If there is more than one GPU, compute operations (other than variable update operations) will be replicated across all GPUs.
In general, when using a multi-worker tf.distribute strategy such as tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy(), there is a tf.distribute.cluster_resolver.ClusterResolver associated with the strategy used, and such an instance is returned by this property.
Strategies that intend to have an associated tf.distribute.cluster_resolver.ClusterResolver must set the relevant attribute, or override this property; otherwise, None is returned by default. Those strategies should also provide information regarding what is returned by this property.
Single-worker strategies usually do not have a tf.distribute.cluster_resolver.ClusterResolver, and in those cases this property will return None.
The tf.distribute.cluster_resolver.ClusterResolver may be useful when the user needs to access information such as the cluster spec, task type or task id. For example,
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's API docstring.
View source
Distributes tf.data.Dataset instances created by calls to dataset_fn.
The argument dataset_fn that users pass in is an input function that has a tf.distribute.InputContext argument and returns a tf.data.Dataset instance. It is expected that the returned dataset from dataset_fn is already batched by per-replica batch size (i.e. global batch size divided by the number of replicas in sync) and sharded. tf.distribute.Strategy.distribute_datasets_from_function does not batch or shard the tf.data.Dataset instance returned from the input function. dataset_fn will be called on the CPU device of each of the workers and each generates a dataset where every replica on that worker will dequeue one batch of inputs (i.e. if a worker has two replicas, two batches will be dequeued from the Dataset every step).
This method can be used for several purposes. First, it allows you to specify your own batching and sharding logic. (In contrast, tf.distribute.experimental_distribute_dataset does batching and sharding for you.) For example, where experimental_distribute_dataset is unable to shard the input files, this method might be used to manually shard the dataset (avoiding the slow fallback behavior in experimental_distribute_dataset). In cases where the dataset is infinite, this sharding can be done by creating dataset replicas that differ only in their random seed.
The dataset_fn should take an tf.distribute.InputContext instance where information about batching and input replication can be accessed.
You can use element_spec property of the tf.distribute.DistributedDataset returned by this API to query the tf.TypeSpec of the elements returned by the iterator. This can be used to set the input_signature property of a tf.function. Follow tf.distribute.DistributedDataset.element_spec to see an example.
For a tutorial on more usage and properties of this method, refer to the tutorial on distributed input). If you are interested in last partial batch handling, read this section.
View source
Distributes a tf.data.Dataset instance provided via dataset.
The returned dataset is a wrapped strategy dataset which creates a multidevice iterator under the hood. It prefetches the input data to the specified devices on the worker. The returned distributed dataset can be iterated over similar to how regular datasets can.
Args: dataset: tf.data.Dataset to be prefetched to device. options: tf.distribute.InputOptions used to control options on how this dataset is distributed.
View source
Generates tf.distribute.DistributedValues from value_fn.
This function is to generate tf.distribute.DistributedValues to pass into run, reduce, or other methods that take distributed values when not using datasets.
View source
Returns the list of all local per-replica values contained in value.
In CentralStorageStrategy there is a single worker so the value returned will be all the values on that worker.
View source
Gather value across replicas along axis to the current device.
Given a tf.distribute.DistributedValues or tf.Tensor-like object value, this API gathers and concatenates value across replicas along the axis-th dimension. The result is copied to the "current" device
This API can only be called in the cross-replica context. For a counterpart in the replica context, see tf.distribute.ReplicaContext.all_gather.
Consider the following example for more combinations:
View source
Reduce value across replicas.
Given a per-replica value returned by run, say a per-example loss, the batch will be divided across all the replicas. This function allows you to aggregate across replicas and optionally also across batch elements. For example, if you have a global batch size of 8 and 2 replicas, values for examples [0, 1, 2, 3] will be on replica 0 and [4, 5, 6, 7] will be on replica 1. By default, reduce will just aggregate across replicas, returning [0+4, 1+5, 2+6, 3+7]. This is useful when each replica is computing a scalar or some other value that doesn't have a "batch" dimension (like a gradient). More often you will want to aggregate across the global batch, which you can get by specifying the batch dimension as the axis, typically axis=0. In this case it would return a scalar 0+1+2+3+4+5+6+7.
If there is a last partial batch, you will need to specify an axis so that the resulting shape is consistent across replicas. So if the last batch has size 6 and it is divided into [0, 1, 2, 3] and [4, 5], you would get a shape mismatch unless you specify axis=0. If you specify tf.distribute.ReduceOp.MEAN, using axis=0 will use the correct denominator of 6. Contrast this with computing reduce_mean to get a scalar value on each replica and this function to average those means, which will weigh some values 1/8 and others 1/4.
View source
Run fn on each replica, with the given arguments.
In CentralStorageStrategy, fn is called on each of the compute replicas, with the provided "per replica" arguments specific to that device.
View source
Context manager to make the strategy current and distribute variables.
This method returns a context manager, and is used as follows:
What happens when Strategy.scope is entered?
What should be in scope and what should be outside?
There are a number of requirements on what needs to happen inside the scope. However, in places where we have information about which strategy is in use, we often enter the scope for the user, so they don't have to do it explicitly (i.e. calling those either inside or outside the scope is OK).


Hints for collective operations like AllReduce.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.experimental.CollectiveHints
This can be passed to methods like tf.distribute.get_replica_context().all_reduce() to optimize collective operation performance. Note that these are only hints, which may or may not change the actual behavior. Some options only apply to certain strategy and are ignored by others.
One common optimization is to break gradients all-reduce into multiple packs so that weight updates can overlap with gradient all-reduce.


Cross device communication implementation.
Main aliases
tf.distribute.experimental.CollectiveCommunication
See Migration guide for more details.
tf.compat.v1.distribute.experimental.CollectiveCommunication, tf.compat.v1.distribute.experimental.CommunicationImplementation


Options for cross device communications like All-reduce.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.distribute.experimental.CommunicationOptions
This can be passed to methods like tf.distribute.get_replica_context().all_reduce() to optimize collective operation performance. Note that these are only hints, which may or may not change the actual behavior. Some options only apply to certain strategy and are ignored by others.
One common optimization is to break gradients all-reduce into multiple packs so that weight updates can overlap with gradient all-reduce.


Public API for tf.distribute.experimental.coordinator namespace.
class ClusterCoordinator: An object to schedule and coordinate remote function execution.
class PerWorkerValues: A container that holds a list of values, one value per worker.
class RemoteValue: An asynchronously available value of a scheduled function.