TPU distribution strategy implementation.
Inherits From: Strategy
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
Creates tf.distribute.DistributedDataset from tf.data.Dataset.
The returned tf.distribute.DistributedDataset can be iterated over similar to regular datasets. NOTE: The user cannot add any more transformations to a tf.distribute.DistributedDataset. You can only create an iterator or examine the tf.TypeSpec of the data generated by it. See API docs of tf.distribute.DistributedDataset to learn more.
The following is an example:
Three key actions happending under the hood of this method are batching, sharding, and prefetching.
In the code snippet above, dataset is batched by global_batch_size, and calling experimental_distribute_dataset on it rebatches dataset to a new batch size that is equal to the global batch size divided by the number of replicas in sync. We iterate through it using a Pythonic for loop. x is a tf.distribute.DistributedValues containing data for all replicas, and each replica gets data of the new batch size. tf.distribute.Strategy.run will take care of feeding the right per-replica data in x to the right replica_fn executed on each replica.
Sharding contains autosharding across multiple workers and within every worker. First, in multi-worker distributed training (i.e. when you use tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy), autosharding a dataset over a set of workers means that each worker is assigned a subset of the entire dataset (if the right tf.data.experimental.AutoShardPolicy is set). This is to ensure that at each step, a global batch size of non-overlapping dataset elements will be processed by each worker. Autosharding has a couple of different options that can be specified using tf.data.experimental.DistributeOptions. Then, sharding within each worker means the method will split the data among all the worker devices (if more than one a present). This will happen regardless of multi-worker autosharding.
By default, this method adds a prefetch transformation at the end of the user provided tf.data.Dataset instance. The argument to the prefetch transformation which is buffer_size is equal to the number of replicas in sync.
If the above batch splitting and dataset sharding logic is undesirable, please use tf.distribute.Strategy.distribute_datasets_from_function instead, which does not do any automatic batching or sharding for you.
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
Returns the list of all local per-replica values contained in value.
View source
Makes a tf.data.Dataset for input provided via a numpy array.
This avoids adding numpy_input as a large constant in the graph, and copies the data to the machine or machines that will be processing the input.
Note that you will likely need to use tf.distribute.Strategy.experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.
View source
Runs ops in fn on each replica, with inputs from input_iterator.
DEPRECATED: This method is not available in TF 2.x. Please switch to using run instead.
When eager execution is enabled, executes ops specified by fn on each replica. Otherwise, builds a graph to execute the ops on each replica.
Each replica will take a single, different input from the inputs provided by one get_next call on the input iterator.
fn may call tf.distribute.get_replica_context() to access members such as replica_id_in_sync_group.
View source
Makes an iterator for input provided via dataset.
DEPRECATED: This method is not available in TF 2.x.
Data from the given dataset will be distributed evenly across all the compute replicas. We will assume that the input dataset is batched by the global batch size. With this assumption, we will make a best effort to divide each batch across all the replicas (one or more workers). If this effort fails, an error will be thrown, and the user should instead use make_input_fn_iterator which provides more control to the user, and does not try to divide a batch across replicas.
The user could also use make_input_fn_iterator if they want to customize which input is fed to which replica/worker etc.
View source
Returns an iterator split across replicas created from an input function.
DEPRECATED: This method is not available in TF 2.x.
The input_fn should take an tf.distribute.InputContext object where information about batching and input sharding can be accessed:
The tf.data.Dataset returned by input_fn should have a per-replica batch size, which may be computed using input_context.get_per_replica_batch_size.
View source
Reduce value across replicas and return result on current device.
To see how this would look with multiple replicas, consider the same example with MirroredStrategy with 2 GPUs:
This API is typically used for aggregating the results returned from different replicas, for reporting etc. For example, loss computed from different replicas can be averaged using this API before printing.
There are a number of different tf.distribute APIs for reducing values across replicas:
What should axis be?
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 by specifying the axis parameter accordingly.
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. With axis=None, reduce will aggregate only 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 or loss).
Sometimes, you will want to aggregate across both the global batch and all replicas. You can get this behavior 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.
Executes ops specified by fn on each replica. If args or kwargs have "per-replica" values, such as those produced by a "distributed Dataset", when fn is executed on a particular replica, it will be executed with the component of those "per-replica" values that correspond to that replica.
fn may call tf.distribute.get_replica_context() to access members such as all_reduce.
All arguments in args or kwargs should either be nest of tensors or per-replica objects containing tensors or composite tensors.
Users can pass strategy specific options to options argument. An example to enable bucketizing dynamic shapes in TPUStrategy.run is:
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).
View source
Returns a copy of config_proto modified for use with this strategy.
DEPRECATED: This method is not available in TF 2.x.
The updated config has something needed to run a strategy, e.g. configuration to run collective ops, or device filters to improve distributed training performance.


tf.distribute.ReduceOp corresponding to the last loss reduction.
This is used to decide whether loss should be scaled in optimizer (used only for estimator + v1 optimizer use case).


Synchronous training across multiple replicas on one machine.
Inherits From: Strategy
This strategy is typically used for training on one machine with multiple GPUs. For TPUs, use tf.distribute.TPUStrategy. To use MirroredStrategy with multiple workers, please refer to tf.distribute.experimental.MultiWorkerMirroredStrategy.
For example, a variable created under a MirroredStrategy is a MirroredVariable. If no devices are specified in the constructor argument of the strategy then it will use all the available GPUs. If no GPUs are found, it will use the available CPUs. Note that TensorFlow treats all CPUs on a machine as a single device, and uses threads internally for parallelism.
While using distribution strategies, all the variable creation should be done within the strategy's scope. This will replicate the variables across all the replicas and keep them in sync using an all-reduce algorithm.
Variables created inside a MirroredStrategy which is wrapped with a tf.function are still MirroredVariables.
experimental_distribute_dataset can be used to distribute the dataset across the replicas when writing your own training loop. If you are using .fit and .compile methods available in tf.keras, then tf.keras will handle the distribution for you.
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
Creates tf.distribute.DistributedDataset from tf.data.Dataset.
The returned tf.distribute.DistributedDataset can be iterated over similar to regular datasets. NOTE: The user cannot add any more transformations to a tf.distribute.DistributedDataset. You can only create an iterator or examine the tf.TypeSpec of the data generated by it. See API docs of tf.distribute.DistributedDataset to learn more.
The following is an example:
Three key actions happending under the hood of this method are batching, sharding, and prefetching.
In the code snippet above, dataset is batched by global_batch_size, and calling experimental_distribute_dataset on it rebatches dataset to a new batch size that is equal to the global batch size divided by the number of replicas in sync. We iterate through it using a Pythonic for loop. x is a tf.distribute.DistributedValues containing data for all replicas, and each replica gets data of the new batch size. tf.distribute.Strategy.run will take care of feeding the right per-replica data in x to the right replica_fn executed on each replica.
Sharding contains autosharding across multiple workers and within every worker. First, in multi-worker distributed training (i.e. when you use tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy), autosharding a dataset over a set of workers means that each worker is assigned a subset of the entire dataset (if the right tf.data.experimental.AutoShardPolicy is set). This is to ensure that at each step, a global batch size of non-overlapping dataset elements will be processed by each worker. Autosharding has a couple of different options that can be specified using tf.data.experimental.DistributeOptions. Then, sharding within each worker means the method will split the data among all the worker devices (if more than one a present). This will happen regardless of multi-worker autosharding.
By default, this method adds a prefetch transformation at the end of the user provided tf.data.Dataset instance. The argument to the prefetch transformation which is buffer_size is equal to the number of replicas in sync.
If the above batch splitting and dataset sharding logic is undesirable, please use tf.distribute.Strategy.distribute_datasets_from_function instead, which does not do any automatic batching or sharding for you.
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
Returns the list of all local per-replica values contained in value.
View source
Makes a tf.data.Dataset for input provided via a numpy array.
This avoids adding numpy_input as a large constant in the graph, and copies the data to the machine or machines that will be processing the input.
Note that you will likely need to use tf.distribute.Strategy.experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.
View source
Runs ops in fn on each replica, with inputs from input_iterator.
DEPRECATED: This method is not available in TF 2.x. Please switch to using run instead.
When eager execution is enabled, executes ops specified by fn on each replica. Otherwise, builds a graph to execute the ops on each replica.
Each replica will take a single, different input from the inputs provided by one get_next call on the input iterator.
fn may call tf.distribute.get_replica_context() to access members such as replica_id_in_sync_group.
View source
Makes an iterator for input provided via dataset.
DEPRECATED: This method is not available in TF 2.x.
Data from the given dataset will be distributed evenly across all the compute replicas. We will assume that the input dataset is batched by the global batch size. With this assumption, we will make a best effort to divide each batch across all the replicas (one or more workers). If this effort fails, an error will be thrown, and the user should instead use make_input_fn_iterator which provides more control to the user, and does not try to divide a batch across replicas.
The user could also use make_input_fn_iterator if they want to customize which input is fed to which replica/worker etc.
View source
Returns an iterator split across replicas created from an input function.
DEPRECATED: This method is not available in TF 2.x.
The input_fn should take an tf.distribute.InputContext object where information about batching and input sharding can be accessed:
The tf.data.Dataset returned by input_fn should have a per-replica batch size, which may be computed using input_context.get_per_replica_batch_size.
View source
Reduce value across replicas and return result on current device.
To see how this would look with multiple replicas, consider the same example with MirroredStrategy with 2 GPUs:
This API is typically used for aggregating the results returned from different replicas, for reporting etc. For example, loss computed from different replicas can be averaged using this API before printing.
There are a number of different tf.distribute APIs for reducing values across replicas:
What should axis be?
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 by specifying the axis parameter accordingly.
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. With axis=None, reduce will aggregate only 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 or loss).
Sometimes, you will want to aggregate across both the global batch and all replicas. You can get this behavior 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
Invokes fn on each replica, with the given arguments.
This method is the primary way to distribute your computation with a tf.distribute object. It invokes fn on each replica. If args or kwargs have tf.distribute.DistributedValues, such as those produced by a tf.distribute.DistributedDataset from tf.distribute.Strategy.experimental_distribute_dataset or tf.distribute.Strategy.distribute_datasets_from_function, when fn is executed on a particular replica, it will be executed with the component of tf.distribute.DistributedValues that correspond to that replica.
fn is invoked under a replica context. fn may call tf.distribute.get_replica_context() to access members such as all_reduce. Please see the module-level docstring of tf.distribute for the concept of replica context.
All arguments in args or kwargs should either be Python values of a nested structure of tensors, e.g. a list of tensors, in which case args and kwargs will be passed to the fn invoked on each replica. Or args or kwargs can be tf.distribute.DistributedValues containing tensors or composite tensors, i.e. tf.compat.v1.TensorInfo.CompositeTensor, in which case each fn call will get the component of a tf.distribute.DistributedValues corresponding to its replica.
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).
View source
Returns a copy of config_proto modified for use with this strategy.
DEPRECATED: This method is not available in TF 2.x.
The updated config has something needed to run a strategy, e.g. configuration to run collective ops, or device filters to improve distributed training performance.


A distribution strategy for running on a single device.
Inherits From: Strategy
Using this strategy will place any variables created in its scope on the specified device. Input distributed through this strategy will be prefetched to the specified device. Moreover, any functions called via strategy.run will also be placed on the specified device as well.
Typical usage of this strategy could be testing your code with the tf.distribute.Strategy API before switching to other strategies which actually distribute to multiple devices/machines.
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
Creates tf.distribute.DistributedDataset from tf.data.Dataset.
The returned tf.distribute.DistributedDataset can be iterated over similar to regular datasets. NOTE: The user cannot add any more transformations to a tf.distribute.DistributedDataset. You can only create an iterator or examine the tf.TypeSpec of the data generated by it. See API docs of tf.distribute.DistributedDataset to learn more.
The following is an example:
Three key actions happending under the hood of this method are batching, sharding, and prefetching.
In the code snippet above, dataset is batched by global_batch_size, and calling experimental_distribute_dataset on it rebatches dataset to a new batch size that is equal to the global batch size divided by the number of replicas in sync. We iterate through it using a Pythonic for loop. x is a tf.distribute.DistributedValues containing data for all replicas, and each replica gets data of the new batch size. tf.distribute.Strategy.run will take care of feeding the right per-replica data in x to the right replica_fn executed on each replica.
Sharding contains autosharding across multiple workers and within every worker. First, in multi-worker distributed training (i.e. when you use tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy), autosharding a dataset over a set of workers means that each worker is assigned a subset of the entire dataset (if the right tf.data.experimental.AutoShardPolicy is set). This is to ensure that at each step, a global batch size of non-overlapping dataset elements will be processed by each worker. Autosharding has a couple of different options that can be specified using tf.data.experimental.DistributeOptions. Then, sharding within each worker means the method will split the data among all the worker devices (if more than one a present). This will happen regardless of multi-worker autosharding.
By default, this method adds a prefetch transformation at the end of the user provided tf.data.Dataset instance. The argument to the prefetch transformation which is buffer_size is equal to the number of replicas in sync.
If the above batch splitting and dataset sharding logic is undesirable, please use tf.distribute.Strategy.distribute_datasets_from_function instead, which does not do any automatic batching or sharding for you.
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
Returns the list of all local per-replica values contained in value.
View source
Makes a tf.data.Dataset for input provided via a numpy array.
This avoids adding numpy_input as a large constant in the graph, and copies the data to the machine or machines that will be processing the input.
Note that you will likely need to use tf.distribute.Strategy.experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.
View source
Runs ops in fn on each replica, with inputs from input_iterator.
DEPRECATED: This method is not available in TF 2.x. Please switch to using run instead.
When eager execution is enabled, executes ops specified by fn on each replica. Otherwise, builds a graph to execute the ops on each replica.
Each replica will take a single, different input from the inputs provided by one get_next call on the input iterator.
fn may call tf.distribute.get_replica_context() to access members such as replica_id_in_sync_group.
View source
Makes an iterator for input provided via dataset.
DEPRECATED: This method is not available in TF 2.x.
Data from the given dataset will be distributed evenly across all the compute replicas. We will assume that the input dataset is batched by the global batch size. With this assumption, we will make a best effort to divide each batch across all the replicas (one or more workers). If this effort fails, an error will be thrown, and the user should instead use make_input_fn_iterator which provides more control to the user, and does not try to divide a batch across replicas.
The user could also use make_input_fn_iterator if they want to customize which input is fed to which replica/worker etc.
View source
Returns an iterator split across replicas created from an input function.
DEPRECATED: This method is not available in TF 2.x.
The input_fn should take an tf.distribute.InputContext object where information about batching and input sharding can be accessed:
The tf.data.Dataset returned by input_fn should have a per-replica batch size, which may be computed using input_context.get_per_replica_batch_size.
View source
Reduce value across replicas and return result on current device.
To see how this would look with multiple replicas, consider the same example with MirroredStrategy with 2 GPUs:
This API is typically used for aggregating the results returned from different replicas, for reporting etc. For example, loss computed from different replicas can be averaged using this API before printing.
There are a number of different tf.distribute APIs for reducing values across replicas:
What should axis be?
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 by specifying the axis parameter accordingly.
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. With axis=None, reduce will aggregate only 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 or loss).
Sometimes, you will want to aggregate across both the global batch and all replicas. You can get this behavior 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
Invokes fn on each replica, with the given arguments.
This method is the primary way to distribute your computation with a tf.distribute object. It invokes fn on each replica. If args or kwargs have tf.distribute.DistributedValues, such as those produced by a tf.distribute.DistributedDataset from tf.distribute.Strategy.experimental_distribute_dataset or tf.distribute.Strategy.distribute_datasets_from_function, when fn is executed on a particular replica, it will be executed with the component of tf.distribute.DistributedValues that correspond to that replica.
fn is invoked under a replica context. fn may call tf.distribute.get_replica_context() to access members such as all_reduce. Please see the module-level docstring of tf.distribute for the concept of replica context.
All arguments in args or kwargs should either be Python values of a nested structure of tensors, e.g. a list of tensors, in which case args and kwargs will be passed to the fn invoked on each replica. Or args or kwargs can be tf.distribute.DistributedValues containing tensors or composite tensors, i.e. tf.compat.v1.TensorInfo.CompositeTensor, in which case each fn call will get the component of a tf.distribute.DistributedValues corresponding to its replica.
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).
View source
Returns a copy of config_proto modified for use with this strategy.
DEPRECATED: This method is not available in TF 2.x.
The updated config has something needed to run a strategy, e.g. configuration to run collective ops, or device filters to improve distributed training performance.


A class with a collection of APIs that can be called in a replica context.
You can use tf.distribute.get_replica_context to get an instance of ReplicaContext, which can only be called inside the function passed to tf.distribute.Strategy.run.
This identifies the replica among all replicas that are kept in sync. The value of the replica id can range from 0 to tf.distribute.ReplicaContext.num_replicas_in_sync - 1.
View source
All-reduces value across all replicas.
It supports batched operations. You can pass a list of values and it attempts to batch them when possible. You can also specify options to indicate the desired batching behavior, e.g. batch the values into multiple packs so that they can better overlap with computations.
Note that all replicas need to participate in the all-reduce, otherwise this operation hangs. Note that if there're multiple all-reduces, they need to execute in the same order on all replicas. Dispatching all-reduce based on conditions is usually error-prone.
This API currently can only be called in the replica context. Other variants to reduce values across replicas are:
View source
Merge args across replicas and run merge_fn in a cross-replica context.
This allows communication and coordination when there are multiple calls to the step_fn triggered by a call to strategy.run(step_fn, ...).
See tf.distribute.Strategy.run for an explanation.
If not inside a distributed scope, this is equivalent to:


A list of devices with a state & compute distribution policy.
See the guide for overview and examples.
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
Creates tf.distribute.DistributedDataset from tf.data.Dataset.
The returned tf.distribute.DistributedDataset can be iterated over similar to regular datasets. NOTE: The user cannot add any more transformations to a tf.distribute.DistributedDataset. You can only create an iterator or examine the tf.TypeSpec of the data generated by it. See API docs of tf.distribute.DistributedDataset to learn more.
The following is an example:
Three key actions happending under the hood of this method are batching, sharding, and prefetching.
In the code snippet above, dataset is batched by global_batch_size, and calling experimental_distribute_dataset on it rebatches dataset to a new batch size that is equal to the global batch size divided by the number of replicas in sync. We iterate through it using a Pythonic for loop. x is a tf.distribute.DistributedValues containing data for all replicas, and each replica gets data of the new batch size. tf.distribute.Strategy.run will take care of feeding the right per-replica data in x to the right replica_fn executed on each replica.
Sharding contains autosharding across multiple workers and within every worker. First, in multi-worker distributed training (i.e. when you use tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy), autosharding a dataset over a set of workers means that each worker is assigned a subset of the entire dataset (if the right tf.data.experimental.AutoShardPolicy is set). This is to ensure that at each step, a global batch size of non-overlapping dataset elements will be processed by each worker. Autosharding has a couple of different options that can be specified using tf.data.experimental.DistributeOptions. Then, sharding within each worker means the method will split the data among all the worker devices (if more than one a present). This will happen regardless of multi-worker autosharding.
By default, this method adds a prefetch transformation at the end of the user provided tf.data.Dataset instance. The argument to the prefetch transformation which is buffer_size is equal to the number of replicas in sync.
If the above batch splitting and dataset sharding logic is undesirable, please use tf.distribute.Strategy.distribute_datasets_from_function instead, which does not do any automatic batching or sharding for you.
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
Returns the list of all local per-replica values contained in value.
View source
Makes a tf.data.Dataset for input provided via a numpy array.
This avoids adding numpy_input as a large constant in the graph, and copies the data to the machine or machines that will be processing the input.
Note that you will likely need to use tf.distribute.Strategy.experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.
View source
Runs ops in fn on each replica, with inputs from input_iterator.
DEPRECATED: This method is not available in TF 2.x. Please switch to using run instead.
When eager execution is enabled, executes ops specified by fn on each replica. Otherwise, builds a graph to execute the ops on each replica.
Each replica will take a single, different input from the inputs provided by one get_next call on the input iterator.
fn may call tf.distribute.get_replica_context() to access members such as replica_id_in_sync_group.
View source
Makes an iterator for input provided via dataset.
DEPRECATED: This method is not available in TF 2.x.
Data from the given dataset will be distributed evenly across all the compute replicas. We will assume that the input dataset is batched by the global batch size. With this assumption, we will make a best effort to divide each batch across all the replicas (one or more workers). If this effort fails, an error will be thrown, and the user should instead use make_input_fn_iterator which provides more control to the user, and does not try to divide a batch across replicas.
The user could also use make_input_fn_iterator if they want to customize which input is fed to which replica/worker etc.
View source
Returns an iterator split across replicas created from an input function.
DEPRECATED: This method is not available in TF 2.x.
The input_fn should take an tf.distribute.InputContext object where information about batching and input sharding can be accessed:
The tf.data.Dataset returned by input_fn should have a per-replica batch size, which may be computed using input_context.get_per_replica_batch_size.
View source
Reduce value across replicas and return result on current device.
To see how this would look with multiple replicas, consider the same example with MirroredStrategy with 2 GPUs:
This API is typically used for aggregating the results returned from different replicas, for reporting etc. For example, loss computed from different replicas can be averaged using this API before printing.
There are a number of different tf.distribute APIs for reducing values across replicas:
What should axis be?
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 by specifying the axis parameter accordingly.
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. With axis=None, reduce will aggregate only 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 or loss).
Sometimes, you will want to aggregate across both the global batch and all replicas. You can get this behavior 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
Invokes fn on each replica, with the given arguments.
This method is the primary way to distribute your computation with a tf.distribute object. It invokes fn on each replica. If args or kwargs have tf.distribute.DistributedValues, such as those produced by a tf.distribute.DistributedDataset from tf.distribute.Strategy.experimental_distribute_dataset or tf.distribute.Strategy.distribute_datasets_from_function, when fn is executed on a particular replica, it will be executed with the component of tf.distribute.DistributedValues that correspond to that replica.
fn is invoked under a replica context. fn may call tf.distribute.get_replica_context() to access members such as all_reduce. Please see the module-level docstring of tf.distribute for the concept of replica context.
All arguments in args or kwargs should either be Python values of a nested structure of tensors, e.g. a list of tensors, in which case args and kwargs will be passed to the fn invoked on each replica. Or args or kwargs can be tf.distribute.DistributedValues containing tensors or composite tensors, i.e. tf.compat.v1.TensorInfo.CompositeTensor, in which case each fn call will get the component of a tf.distribute.DistributedValues corresponding to its replica.
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).
View source
Returns a copy of config_proto modified for use with this strategy.
DEPRECATED: This method is not available in TF 2.x.
The updated config has something needed to run a strategy, e.g. configuration to run collective ops, or device filters to improve distributed training performance.


Additional APIs for algorithms that need to be distribution-aware.
Inherits From: StrategyExtended
Some common use cases of functions on this page:
tf.distribute.DistributedValues can have the same locality as a distributed variable, which leads to a mirrored value residing on the same devices as the variable (as opposed to the compute devices). Such values may be passed to a call to tf.distribute.StrategyExtended.update to update the value of a variable. You may use tf.distribute.StrategyExtended.colocate_vars_with to give a variable the same locality as another variable. You may convert a "PerReplica" value to a variable's locality by using tf.distribute.StrategyExtended.reduce_to or tf.distribute.StrategyExtended.batch_reduce_to.
A distributed variable is variables created on multiple devices. As discussed in the glossary, mirrored variable and SyncOnRead variable are two examples. The standard pattern for updating distributed variables is to:
Steps 2 through 4 are done automatically by class tf.keras.optimizers.Optimizer if you call its tf.keras.optimizers.Optimizer.apply_gradients method in a replica context.
In fact, a higher-level solution to update a distributed variable is by calling assign on the variable as you would do to a regular tf.Variable. You can call the method in both replica context and cross-replica context. For a mirrored variable, calling assign in replica context requires you to specify the aggregation type in the variable constructor. In that case, the context switching and sync described in steps 2 through 4 are handled for you. If you call assign on mirrored variable in cross-replica context, you can only assign a single value or assign values from another mirrored variable or a mirrored tf.distribute.DistributedValues. For a SyncOnRead variable, in replica context, you can simply call assign on it and no aggregation happens under the hood. In cross-replica context, you can only assign a single value to a SyncOnRead variable. One example case is restoring from a checkpoint: if the aggregation type of the variable is tf.VariableAggregation.SUM, it is assumed that replica values were added before checkpointing, so at the time of restoring, the value is divided by the number of replicas and then assigned to each replica; if the aggregation type is tf.VariableAggregation.MEAN, the value is assigned to each replica directly.
This is expected to return a constant value that will not be changed throughout its life cycle.
View source
Combine multiple reduce_to calls into one for faster execution.
Similar to reduce_to, but accepts a list of (value, destinations) pairs. It's more efficient than reduce each value separately.
This API currently can only be called in cross-replica context. Other variants to reduce values across replicas are:
See reduce_to for more information.
View source
Mirror a tensor on one device to all worker devices.
View source
Run fn once per replica.
fn may call tf.get_replica_context() to access methods such as replica_id_in_sync_group and merge_call().
merge_call() is used to communicate between the replicas and re-enter the cross-replica context. All replicas pause their execution having encountered a merge_call() call. After that the merge_fn-function is executed. Its results are then unwrapped and given back to each replica call. After that execution resumes until fn is complete or encounters another merge_call(). Example:
View source
Scope that controls which devices variables will be created on.
No operations should be added to the graph inside this scope, it should only be used when creating variables (some implementations work by changing variable creation, others work by using a tf.compat.v1.colocate_with() scope).
This may only be used inside self.scope().
View source
Makes a dataset for input provided via a numpy array.
This avoids adding numpy_input as a large constant in the graph, and copies the data to the machine or machines that will be processing the input.
View source
DEPRECATED: please use run instead.
Run fn with input from iterator for iterations times.
This method can be used to run a step function for training a number of times using input from a dataset.
View source
Device(s) for non-slot variables.
DEPRECATED: TF 1.x ONLY.
This method returns non-slot devices where non-slot variables are placed. Users can create non-slot variables on these devices by using a block:
View source
Reads the value of a variable.
Returns the aggregate value of a replica-local variable, or the (read-only) value of any other variable.
View source
Combine (via e.g. sum or mean) values across replicas.
reduce_to aggregates tf.distribute.DistributedValues and distributed variables. It supports both dense values and tf.IndexedSlices.
This API currently can only be called in cross-replica context. Other variants to reduce values across replicas are:
destinations specifies where to reduce the value to, e.g. "GPU:0". You can also pass in a Tensor, and the destinations will be the device of that tensor. For all-reduce, pass the same to value and destinations.
It can be used in tf.distribute.ReplicaContext.merge_call to write code that works for all tf.distribute.Strategy.
View source
Run fn to update var using inputs mirrored to the same devices.
tf.distribute.StrategyExtended.update takes a distributed variable var to be updated, an update function fn, and args and kwargs for fn. It applies fn to each component variable of var and passes corresponding values from args and kwargs. Neither args nor kwargs may contain per-replica values. If they contain mirrored values, they will be unwrapped before calling fn. For example, fn can be assign_add and args can be a mirrored DistributedValues where each component contains the value to be added to this mirrored variable var. Calling update will call assign_add on each component variable of var with the corresponding tensor value on that device.
If var is mirrored across multiple devices, then this method implements logic as following:
Otherwise, this method returns fn(var, *args, **kwargs) colocated with var.
View source
Runs fn(*args, **kwargs) on colocate_with devices.
Used to update non-slot variables.
DEPRECATED: TF 1.x ONLY.
View source
Returns the container that this per-replica value belongs to.
View source
Tests whether v was created while this strategy scope was active.
Variables created inside the strategy scope are "owned" by it:
Variables created outside the strategy are not owned by it:


Core module for TensorFlow distribution objects and helpers.
class Bernoulli: Bernoulli distribution.
class Beta: Beta distribution.
class Categorical: Categorical distribution.
class Dirichlet: Dirichlet distribution.
class DirichletMultinomial: Dirichlet-Multinomial compound distribution.
class Distribution: A generic probability distribution base class.
class Exponential: Exponential distribution.
class Gamma: Gamma distribution.
class Laplace: The Laplace distribution with location loc and scale parameters.
class Multinomial: Multinomial distribution.
class Normal: The Normal distribution with location loc and scale parameters.
class RegisterKL: Decorator to register a KL divergence implementation function.
class ReparameterizationType: Instances of this class represent how sampling is reparameterized.
class StudentT: Student's t-distribution.
class Uniform: Uniform distribution with low and high parameters.
kl_divergence(...): Get the KL-divergence KL(distribution_a || distribution_b). (deprecated)


Bernoulli distribution.
Inherits From: Distribution
The Bernoulli distribution with probs parameter, i.e., the probability of a 1 outcome (vs a 0 outcome).
Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
May be partially defined or unknown.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
May be partially defined or unknown.
Currently this is one of the static instances distributions.FULLY_REPARAMETERIZED or distributions.NOT_REPARAMETERIZED.
View source
Shape of a single sample from a single event index as a 1-D Tensor.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
View source
Cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
View source
Creates a deep copy of the distribution.
View source
Covariance.
Covariance is (possibly) defined only for non-scalar-event distributions.
For example, for a length-k, vector-valued distribution, it is calculated as,
where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.
Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,
where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.
View source
Computes the (Shannon) cross entropy.
Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shanon) cross entropy is defined as:
where F denotes the support of the random variable X ~ P.
View source
Shannon entropy in nats.
View source
Shape of a single sample from a single batch as a 1-D int32 Tensor.
View source
Indicates that batch_shape == [].
View source
Indicates that event_shape == [].
View source
Computes the Kullback--Leibler divergence.
Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:
where F denotes the support of the random variable X ~ p, H[., .] denotes (Shanon) cross entropy, and H[.] denotes (Shanon) entropy.
View source
Log cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.
View source
Log probability density/mass function.
View source
Log survival function.
Given random variable X, the survival function is defined:
Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.
View source
Mean.
View source
Mode.
Additional documentation from Bernoulli:
Returns 1 if prob > 0.5 and 0 otherwise.
View source
Shapes of parameters given the desired shape of a call to sample().
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().
Subclasses should override class method _param_shapes.
View source
param_shapes with static (i.e. TensorShape) shapes.
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.
Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.
View source
Probability density/mass function.
View source
Quantile function. Aka "inverse cdf" or "percent point function".
Given random variable X and p in [0, 1], the quantile is:
View source
Generate samples of the specified shape.
Note that a call to sample() without arguments will generate a single sample.
View source
Standard deviation.
Standard deviation is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.
View source
Survival function.
Given random variable X, the survival function is defined:
View source
Variance.
Variance is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.


Beta distribution.
Inherits From: Distribution
The Beta distribution is defined over the (0, 1) interval using parameters concentration1 (aka "alpha") and concentration0 (aka "beta").
The probability density function (pdf) is,
where:
The concentration parameters represent mean total counts of a 1 or a 0, i.e.,
where mean in (0, 1) and total_concentration is a positive real number representing a mean total_count = concentration1 + concentration0.
Distribution parameters are automatically broadcast in all functions; see examples for details.
Samples of this distribution are reparameterized (pathwise differentiable). The derivatives are computed using the approach described in (Figurnov et al., 2018).
Compute the gradients of samples w.r.t. the parameters:
Implicit Reparameterization Gradients: Figurnov et al., 2018 (pdf)
Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
May be partially defined or unknown.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
May be partially defined or unknown.
Currently this is one of the static instances distributions.FULLY_REPARAMETERIZED or distributions.NOT_REPARAMETERIZED.
View source
Shape of a single sample from a single event index as a 1-D Tensor.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
View source
Cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Additional documentation from Beta:
View source
Creates a deep copy of the distribution.
View source
Covariance.
Covariance is (possibly) defined only for non-scalar-event distributions.
For example, for a length-k, vector-valued distribution, it is calculated as,
where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.
Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,
where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.
View source
Computes the (Shannon) cross entropy.
Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shanon) cross entropy is defined as:
where F denotes the support of the random variable X ~ P.
View source
Shannon entropy in nats.
View source
Shape of a single sample from a single batch as a 1-D int32 Tensor.
View source
Indicates that batch_shape == [].
View source
Indicates that event_shape == [].
View source
Computes the Kullback--Leibler divergence.
Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:
where F denotes the support of the random variable X ~ p, H[., .] denotes (Shanon) cross entropy, and H[.] denotes (Shanon) entropy.
View source
Log cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.
Additional documentation from Beta:
View source
Log probability density/mass function.
Additional documentation from Beta:
View source
Log survival function.
Given random variable X, the survival function is defined:
Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.
View source
Mean.
View source
Mode.
Additional documentation from Beta:
View source
Shapes of parameters given the desired shape of a call to sample().
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().
Subclasses should override class method _param_shapes.
View source
param_shapes with static (i.e. TensorShape) shapes.
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.
Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.
View source
Probability density/mass function.
Additional documentation from Beta:
View source
Quantile function. Aka "inverse cdf" or "percent point function".
Given random variable X and p in [0, 1], the quantile is:
View source
Generate samples of the specified shape.
Note that a call to sample() without arguments will generate a single sample.
View source
Standard deviation.
Standard deviation is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.
View source
Survival function.
Given random variable X, the survival function is defined:
View source
Variance.
Variance is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.


Categorical distribution.
Inherits From: Distribution
The Categorical distribution is parameterized by either probabilities or log-probabilities of a set of K classes. It is defined over the integers {0, 1, ..., K}.
The Categorical distribution is closely related to the OneHotCategorical and Multinomial distributions. The Categorical distribution can be intuited as generating samples according to argmax{ OneHotCategorical(probs) } itself being identical to argmax{ Multinomial(probs, total_count=1) }.
The probability mass function (pmf) is,
The number of classes, K, must not exceed:
In other words,
Creates a 3-class distribution with the 2nd class being most likely.
Creates a 3-class distribution with the 2nd class being most likely. Parameterized by logits rather than probabilities.
Creates a 3-class distribution with the 3rd class being most likely. The distribution functions can be evaluated on counts.
Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
May be partially defined or unknown.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
May be partially defined or unknown.
Currently this is one of the static instances distributions.FULLY_REPARAMETERIZED or distributions.NOT_REPARAMETERIZED.
View source
Shape of a single sample from a single event index as a 1-D Tensor.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
View source
Cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
View source
Creates a deep copy of the distribution.
View source
Covariance.
Covariance is (possibly) defined only for non-scalar-event distributions.
For example, for a length-k, vector-valued distribution, it is calculated as,
where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.
Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,
where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.
View source
Computes the (Shannon) cross entropy.
Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shanon) cross entropy is defined as:
where F denotes the support of the random variable X ~ P.
View source
Shannon entropy in nats.
View source
Shape of a single sample from a single batch as a 1-D int32 Tensor.
View source
Indicates that batch_shape == [].
View source
Indicates that event_shape == [].
View source
Computes the Kullback--Leibler divergence.
Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:
where F denotes the support of the random variable X ~ p, H[., .] denotes (Shanon) cross entropy, and H[.] denotes (Shanon) entropy.
View source
Log cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.
View source
Log probability density/mass function.
View source
Log survival function.
Given random variable X, the survival function is defined:
Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.
View source
Mean.
View source
Mode.
View source
Shapes of parameters given the desired shape of a call to sample().
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().
Subclasses should override class method _param_shapes.
View source
param_shapes with static (i.e. TensorShape) shapes.
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.
Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.
View source
Probability density/mass function.
View source
Quantile function. Aka "inverse cdf" or "percent point function".
Given random variable X and p in [0, 1], the quantile is:
View source
Generate samples of the specified shape.
Note that a call to sample() without arguments will generate a single sample.
View source
Standard deviation.
Standard deviation is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.
View source
Survival function.
Given random variable X, the survival function is defined:
View source
Variance.
Variance is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.


Dirichlet distribution.
Inherits From: Distribution
The Dirichlet distribution is defined over the (k-1)-simplex using a positive, length-k vector concentration (k > 1). The Dirichlet is identically the Beta distribution when k = 2.
The Dirichlet is a distribution over the open (k-1)-simplex, i.e.,
The probability density function (pdf) is,
where:
The concentration represents mean total counts of class occurrence, i.e.,
where mean in S^{k-1} and total_concentration is a positive real number representing a mean total count.
Distribution parameters are automatically broadcast in all functions; see examples for details.
Samples of this distribution are reparameterized (pathwise differentiable). The derivatives are computed using the approach described in (Figurnov et al., 2018).
Compute the gradients of samples w.r.t. the parameters:
Implicit Reparameterization Gradients: Figurnov et al., 2018 (pdf)
Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
May be partially defined or unknown.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
May be partially defined or unknown.
Currently this is one of the static instances distributions.FULLY_REPARAMETERIZED or distributions.NOT_REPARAMETERIZED.
View source
Shape of a single sample from a single event index as a 1-D Tensor.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
View source
Cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
View source
Creates a deep copy of the distribution.
View source
Covariance.
Covariance is (possibly) defined only for non-scalar-event distributions.
For example, for a length-k, vector-valued distribution, it is calculated as,
where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.
Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,
where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.
View source
Computes the (Shannon) cross entropy.
Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shanon) cross entropy is defined as:
where F denotes the support of the random variable X ~ P.
View source
Shannon entropy in nats.
View source
Shape of a single sample from a single batch as a 1-D int32 Tensor.
View source
Indicates that batch_shape == [].
View source
Indicates that event_shape == [].
View source
Computes the Kullback--Leibler divergence.
Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:
where F denotes the support of the random variable X ~ p, H[., .] denotes (Shanon) cross entropy, and H[.] denotes (Shanon) entropy.
View source
Log cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.
View source
Log probability density/mass function.
Additional documentation from Dirichlet:
View source
Log survival function.
Given random variable X, the survival function is defined:
Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.
View source
Mean.
View source
Mode.
Additional documentation from Dirichlet:
View source
Shapes of parameters given the desired shape of a call to sample().
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().
Subclasses should override class method _param_shapes.
View source
param_shapes with static (i.e. TensorShape) shapes.
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.
Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.
View source
Probability density/mass function.
Additional documentation from Dirichlet:
View source
Quantile function. Aka "inverse cdf" or "percent point function".
Given random variable X and p in [0, 1], the quantile is:
View source
Generate samples of the specified shape.
Note that a call to sample() without arguments will generate a single sample.
View source
Standard deviation.
Standard deviation is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.
View source
Survival function.
Given random variable X, the survival function is defined:
View source
Variance.
Variance is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.


Dirichlet-Multinomial compound distribution.
Inherits From: Distribution
The Dirichlet-Multinomial distribution is parameterized by a (batch of) length-K concentration vectors (K > 1) and a total_count number of trials, i.e., the number of trials per draw from the DirichletMultinomial. It is defined over a (batch of) length-K vector counts such that tf.reduce_sum(counts, -1) = total_count. The Dirichlet-Multinomial is identically the Beta-Binomial distribution when K = 2.
The Dirichlet-Multinomial is a distribution over K-class counts, i.e., a length-K vector of non-negative integer counts = n = [n_0, ..., n_{K-1}].
The probability mass function (pmf) is,
where:
Dirichlet-Multinomial is a compound distribution, i.e., its samples are generated as follows.
The last concentration dimension parametrizes a single Dirichlet-Multinomial distribution. When calling distribution functions (e.g., dist.prob(counts)), concentration, total_count and counts are broadcast to the same shape. The last dimension of counts corresponds single Dirichlet-Multinomial distributions.
Distribution parameters are automatically broadcast in all functions; see examples for details.
The number of classes, K, must not exceed:
In other words,
Creates a 3-class distribution, with the 3rd class is most likely to be drawn. The distribution functions can be evaluated on counts.
Creates a 2-batch of 3-class distributions.
Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
May be partially defined or unknown.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
May be partially defined or unknown.
Currently this is one of the static instances distributions.FULLY_REPARAMETERIZED or distributions.NOT_REPARAMETERIZED.
View source
Shape of a single sample from a single event index as a 1-D Tensor.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
View source
Cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
View source
Creates a deep copy of the distribution.
View source
Covariance.
Covariance is (possibly) defined only for non-scalar-event distributions.
For example, for a length-k, vector-valued distribution, it is calculated as,
where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.
Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,
where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.
Additional documentation from DirichletMultinomial:
The covariance for each batch member is defined as the following:
where concentration = alpha and total_concentration = alpha_0 = sum_j alpha_j.
The covariance between elements in a batch is defined as:
View source
Computes the (Shannon) cross entropy.
Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shanon) cross entropy is defined as:
where F denotes the support of the random variable X ~ P.
View source
Shannon entropy in nats.
View source
Shape of a single sample from a single batch as a 1-D int32 Tensor.
View source
Indicates that batch_shape == [].
View source
Indicates that event_shape == [].
View source
Computes the Kullback--Leibler divergence.
Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:
where F denotes the support of the random variable X ~ p, H[., .] denotes (Shanon) cross entropy, and H[.] denotes (Shanon) entropy.
View source
Log cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.
View source
Log probability density/mass function.
Additional documentation from DirichletMultinomial:
For each batch of counts, value = [n_0, ..., n_{K-1}], P[value] is the probability that after sampling self.total_count draws from this Dirichlet-Multinomial distribution, the number of draws falling in class j is n_j. Since this definition is exchangeable; different sequences have the same counts so the probability includes a combinatorial coefficient.
View source
Log survival function.
Given random variable X, the survival function is defined:
Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.
View source
Mean.
View source
Mode.
View source
Shapes of parameters given the desired shape of a call to sample().
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().
Subclasses should override class method _param_shapes.
View source
param_shapes with static (i.e. TensorShape) shapes.
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.
Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.
View source
Probability density/mass function.
Additional documentation from DirichletMultinomial:
For each batch of counts, value = [n_0, ..., n_{K-1}], P[value] is the probability that after sampling self.total_count draws from this Dirichlet-Multinomial distribution, the number of draws falling in class j is n_j. Since this definition is exchangeable; different sequences have the same counts so the probability includes a combinatorial coefficient.
View source
Quantile function. Aka "inverse cdf" or "percent point function".
Given random variable X and p in [0, 1], the quantile is:
View source
Generate samples of the specified shape.
Note that a call to sample() without arguments will generate a single sample.
View source
Standard deviation.
Standard deviation is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.
View source
Survival function.
Given random variable X, the survival function is defined:
View source
Variance.
Variance is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.


A generic probability distribution base class.
Distribution is a base class for constructing and organizing properties (e.g., mean, variance) of random variables (e.g, Bernoulli, Gaussian).
Subclasses are expected to implement a leading-underscore version of the same-named function. The argument signature should be identical except for the omission of name="...". For example, to enable log_prob(value, name="log_prob") a subclass should implement _log_prob(value).
Subclasses can append to public-level docstrings by providing docstrings for their method specializations. For example:
would add the string "Some other details." to the log_prob function docstring. This is implemented as a simple decorator to avoid python linter complaining about missing Args/Returns/Raises sections in the partial docstrings.
All distributions support batches of independent distributions of that type. The batch shape is determined by broadcasting together the parameters.
The shape of arguments to __init__, cdf, log_cdf, prob, and log_prob reflect this broadcasting, as does the return value of sample and sample_n.
sample_n_shape = [n] + batch_shape + event_shape, where sample_n_shape is the shape of the Tensor returned from sample_n, n is the number of samples, batch_shape defines how many independent distributions there are, and event_shape defines the shape of samples from each of those independent distributions. Samples are independent along the batch_shape dimensions, but not necessarily so along the event_shape dimensions (depending on the particulars of the underlying distribution).
Using the Uniform distribution as an example:
There are three important concepts associated with TensorFlow Distributions shapes:
The event shape and the batch shape are properties of a Distribution object, whereas the sample shape is associated with a specific call to sample or log_prob.
For detailed usage examples of TensorFlow Distributions shapes, see this tutorial
Some distributions do not have well-defined statistics for all initialization parameter values. For example, the beta distribution is parameterized by positive real numbers concentration1 and concentration0, and does not have well-defined mode if concentration1 < 1 or concentration0 < 1.
The user is given the option of raising an exception or returning NaN.
In all cases, an exception is raised if invalid parameters are passed, e.g.
Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
May be partially defined or unknown.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
May be partially defined or unknown.
Currently this is one of the static instances distributions.FULLY_REPARAMETERIZED or distributions.NOT_REPARAMETERIZED.
View source
Shape of a single sample from a single event index as a 1-D Tensor.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
View source
Cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
View source
Creates a deep copy of the distribution.
View source
Covariance.
Covariance is (possibly) defined only for non-scalar-event distributions.
For example, for a length-k, vector-valued distribution, it is calculated as,
where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.
Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,
where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.
View source
Computes the (Shannon) cross entropy.
Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shanon) cross entropy is defined as:
where F denotes the support of the random variable X ~ P.
View source
Shannon entropy in nats.
View source
Shape of a single sample from a single batch as a 1-D int32 Tensor.
View source
Indicates that batch_shape == [].
View source
Indicates that event_shape == [].
View source
Computes the Kullback--Leibler divergence.
Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:
where F denotes the support of the random variable X ~ p, H[., .] denotes (Shanon) cross entropy, and H[.] denotes (Shanon) entropy.
View source
Log cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.
View source
Log probability density/mass function.
View source
Log survival function.
Given random variable X, the survival function is defined:
Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.
View source
Mean.
View source
Mode.
View source
Shapes of parameters given the desired shape of a call to sample().
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().
Subclasses should override class method _param_shapes.
View source
param_shapes with static (i.e. TensorShape) shapes.
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.
Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.
View source
Probability density/mass function.
View source
Quantile function. Aka "inverse cdf" or "percent point function".
Given random variable X and p in [0, 1], the quantile is:
View source
Generate samples of the specified shape.
Note that a call to sample() without arguments will generate a single sample.
View source
Standard deviation.
Standard deviation is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.
View source
Survival function.
Given random variable X, the survival function is defined:
View source
Variance.
Variance is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.


Exponential distribution.
Inherits From: Gamma, Distribution
The Exponential distribution is parameterized by an event rate parameter.
The probability density function (pdf) is,
where rate = lambda and Z is the normalizaing constant.
The Exponential distribution is a special case of the Gamma distribution, i.e.,
The Exponential distribution uses a rate parameter, or "inverse scale", which can be intuited as,
Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
May be partially defined or unknown.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
May be partially defined or unknown.
Currently this is one of the static instances distributions.FULLY_REPARAMETERIZED or distributions.NOT_REPARAMETERIZED.
View source
Shape of a single sample from a single event index as a 1-D Tensor.
The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
View source
Cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
View source
Creates a deep copy of the distribution.
View source
Covariance.
Covariance is (possibly) defined only for non-scalar-event distributions.
For example, for a length-k, vector-valued distribution, it is calculated as,
where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.
Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,
where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.
View source
Computes the (Shannon) cross entropy.
Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shanon) cross entropy is defined as:
where F denotes the support of the random variable X ~ P.
View source
Shannon entropy in nats.
View source
Shape of a single sample from a single batch as a 1-D int32 Tensor.
View source
Indicates that batch_shape == [].
View source
Indicates that event_shape == [].
View source
Computes the Kullback--Leibler divergence.
Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:
where F denotes the support of the random variable X ~ p, H[., .] denotes (Shanon) cross entropy, and H[.] denotes (Shanon) entropy.
View source
Log cumulative distribution function.
Given random variable X, the cumulative distribution function cdf is:
Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.
View source
Log probability density/mass function.
View source
Log survival function.
Given random variable X, the survival function is defined:
Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.
View source
Mean.
View source
Mode.
Additional documentation from Gamma:
The mode of a gamma distribution is (shape - 1) / rate when shape > 1, and NaN otherwise. If self.allow_nan_stats is False, an exception will be raised rather than returning NaN.
View source
Shapes of parameters given the desired shape of a call to sample().
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().
Subclasses should override class method _param_shapes.
View source
param_shapes with static (i.e. TensorShape) shapes.
This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.
Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.
View source
Probability density/mass function.
View source
Quantile function. Aka "inverse cdf" or "percent point function".
Given random variable X and p in [0, 1], the quantile is:
View source
Generate samples of the specified shape.
Note that a call to sample() without arguments will generate a single sample.
View source
Standard deviation.
Standard deviation is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.
View source
Survival function.
Given random variable X, the survival function is defined:
View source
Variance.
Variance is defined as,
where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.