| # Copyright 2016 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """The Dirichlet distribution class.""" |
| |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| |
| import numpy as np |
| |
| from tensorflow.python.framework import ops |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import check_ops |
| from tensorflow.python.ops import control_flow_ops |
| from tensorflow.python.ops import math_ops |
| from tensorflow.python.ops import random_ops |
| from tensorflow.python.ops import special_math_ops |
| from tensorflow.python.ops.distributions import distribution |
| from tensorflow.python.ops.distributions import kullback_leibler |
| from tensorflow.python.ops.distributions import util as distribution_util |
| from tensorflow.python.util.tf_export import tf_export |
| |
| |
| __all__ = [ |
| "Dirichlet", |
| ] |
| |
| |
| _dirichlet_sample_note = """Note: `value` must be a non-negative tensor with |
| dtype `self.dtype` and be in the `(self.event_shape() - 1)`-simplex, i.e., |
| `tf.reduce_sum(value, -1) = 1`. It must have a shape compatible with |
| `self.batch_shape() + self.event_shape()`.""" |
| |
| |
| @tf_export("distributions.Dirichlet") |
| class Dirichlet(distribution.Distribution): |
| """Dirichlet distribution. |
| |
| The Dirichlet distribution is defined over the |
| [`(k-1)`-simplex](https://en.wikipedia.org/wiki/Simplex) using a positive, |
| length-`k` vector `concentration` (`k > 1`). The Dirichlet is identically the |
| Beta distribution when `k = 2`. |
| |
| #### Mathematical Details |
| |
| The Dirichlet is a distribution over the open `(k-1)`-simplex, i.e., |
| |
| ```none |
| S^{k-1} = { (x_0, ..., x_{k-1}) in R^k : sum_j x_j = 1 and all_j x_j > 0 }. |
| ``` |
| |
| The probability density function (pdf) is, |
| |
| ```none |
| pdf(x; alpha) = prod_j x_j**(alpha_j - 1) / Z |
| Z = prod_j Gamma(alpha_j) / Gamma(sum_j alpha_j) |
| ``` |
| |
| where: |
| |
| * `x in S^{k-1}`, i.e., the `(k-1)`-simplex, |
| * `concentration = alpha = [alpha_0, ..., alpha_{k-1}]`, `alpha_j > 0`, |
| * `Z` is the normalization constant aka the [multivariate beta function]( |
| https://en.wikipedia.org/wiki/Beta_function#Multivariate_beta_function), |
| and, |
| * `Gamma` is the [gamma function]( |
| https://en.wikipedia.org/wiki/Gamma_function). |
| |
| The `concentration` represents mean total counts of class occurrence, i.e., |
| |
| ```none |
| concentration = alpha = mean * total_concentration |
| ``` |
| |
| 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. |
| |
| Warning: Some components of the samples can be zero due to finite precision. |
| This happens more often when some of the concentrations are very small. |
| Make sure to round the samples to `np.finfo(dtype).tiny` before computing the |
| density. |
| |
| Samples of this distribution are reparameterized (pathwise differentiable). |
| The derivatives are computed using the approach described in the paper |
| |
| [Michael Figurnov, Shakir Mohamed, Andriy Mnih. |
| Implicit Reparameterization Gradients, 2018](https://arxiv.org/abs/1805.08498) |
| |
| #### Examples |
| |
| ```python |
| # Create a single trivariate Dirichlet, with the 3rd class being three times |
| # more frequent than the first. I.e., batch_shape=[], event_shape=[3]. |
| alpha = [1., 2, 3] |
| dist = tf.distributions.Dirichlet(alpha) |
| |
| dist.sample([4, 5]) # shape: [4, 5, 3] |
| |
| # x has one sample, one batch, three classes: |
| x = [.2, .3, .5] # shape: [3] |
| dist.prob(x) # shape: [] |
| |
| # x has two samples from one batch: |
| x = [[.1, .4, .5], |
| [.2, .3, .5]] |
| dist.prob(x) # shape: [2] |
| |
| # alpha will be broadcast to shape [5, 7, 3] to match x. |
| x = [[...]] # shape: [5, 7, 3] |
| dist.prob(x) # shape: [5, 7] |
| ``` |
| |
| ```python |
| # Create batch_shape=[2], event_shape=[3]: |
| alpha = [[1., 2, 3], |
| [4, 5, 6]] # shape: [2, 3] |
| dist = tf.distributions.Dirichlet(alpha) |
| |
| dist.sample([4, 5]) # shape: [4, 5, 2, 3] |
| |
| x = [.2, .3, .5] |
| # x will be broadcast as [[.2, .3, .5], |
| # [.2, .3, .5]], |
| # thus matching batch_shape [2, 3]. |
| dist.prob(x) # shape: [2] |
| ``` |
| |
| Compute the gradients of samples w.r.t. the parameters: |
| |
| ```python |
| alpha = tf.constant([1.0, 2.0, 3.0]) |
| dist = tf.distributions.Dirichlet(alpha) |
| samples = dist.sample(5) # Shape [5, 3] |
| loss = tf.reduce_mean(tf.square(samples)) # Arbitrary loss function |
| # Unbiased stochastic gradients of the loss function |
| grads = tf.gradients(loss, alpha) |
| ``` |
| |
| """ |
| |
| def __init__(self, |
| concentration, |
| validate_args=False, |
| allow_nan_stats=True, |
| name="Dirichlet"): |
| """Initialize a batch of Dirichlet distributions. |
| |
| Args: |
| concentration: Positive floating-point `Tensor` indicating mean number |
| of class occurrences; aka "alpha". Implies `self.dtype`, and |
| `self.batch_shape`, `self.event_shape`, i.e., if |
| `concentration.shape = [N1, N2, ..., Nm, k]` then |
| `batch_shape = [N1, N2, ..., Nm]` and |
| `event_shape = [k]`. |
| validate_args: Python `bool`, default `False`. When `True` distribution |
| parameters are checked for validity despite possibly degrading runtime |
| performance. When `False` invalid inputs may silently render incorrect |
| outputs. |
| allow_nan_stats: Python `bool`, default `True`. When `True`, statistics |
| (e.g., mean, mode, variance) use the value "`NaN`" to indicate the |
| result is undefined. When `False`, an exception is raised if one or |
| more of the statistic's batch members are undefined. |
| name: Python `str` name prefixed to Ops created by this class. |
| """ |
| parameters = dict(locals()) |
| with ops.name_scope(name, values=[concentration]) as name: |
| self._concentration = self._maybe_assert_valid_concentration( |
| ops.convert_to_tensor(concentration, name="concentration"), |
| validate_args) |
| self._total_concentration = math_ops.reduce_sum(self._concentration, -1) |
| super(Dirichlet, self).__init__( |
| dtype=self._concentration.dtype, |
| validate_args=validate_args, |
| allow_nan_stats=allow_nan_stats, |
| reparameterization_type=distribution.FULLY_REPARAMETERIZED, |
| parameters=parameters, |
| graph_parents=[self._concentration, |
| self._total_concentration], |
| name=name) |
| |
| @property |
| def concentration(self): |
| """Concentration parameter; expected counts for that coordinate.""" |
| return self._concentration |
| |
| @property |
| def total_concentration(self): |
| """Sum of last dim of concentration parameter.""" |
| return self._total_concentration |
| |
| def _batch_shape_tensor(self): |
| return array_ops.shape(self.total_concentration) |
| |
| def _batch_shape(self): |
| return self.total_concentration.get_shape() |
| |
| def _event_shape_tensor(self): |
| return array_ops.shape(self.concentration)[-1:] |
| |
| def _event_shape(self): |
| return self.concentration.get_shape().with_rank_at_least(1)[-1:] |
| |
| def _sample_n(self, n, seed=None): |
| gamma_sample = random_ops.random_gamma( |
| shape=[n], |
| alpha=self.concentration, |
| dtype=self.dtype, |
| seed=seed) |
| return gamma_sample / math_ops.reduce_sum(gamma_sample, -1, keepdims=True) |
| |
| @distribution_util.AppendDocstring(_dirichlet_sample_note) |
| def _log_prob(self, x): |
| return self._log_unnormalized_prob(x) - self._log_normalization() |
| |
| @distribution_util.AppendDocstring(_dirichlet_sample_note) |
| def _prob(self, x): |
| return math_ops.exp(self._log_prob(x)) |
| |
| def _log_unnormalized_prob(self, x): |
| x = self._maybe_assert_valid_sample(x) |
| return math_ops.reduce_sum((self.concentration - 1.) * math_ops.log(x), -1) |
| |
| def _log_normalization(self): |
| return special_math_ops.lbeta(self.concentration) |
| |
| def _entropy(self): |
| k = math_ops.cast(self.event_shape_tensor()[0], self.dtype) |
| return ( |
| self._log_normalization() |
| + ((self.total_concentration - k) |
| * math_ops.digamma(self.total_concentration)) |
| - math_ops.reduce_sum( |
| (self.concentration - 1.) * math_ops.digamma(self.concentration), |
| axis=-1)) |
| |
| def _mean(self): |
| return self.concentration / self.total_concentration[..., array_ops.newaxis] |
| |
| def _covariance(self): |
| x = self._variance_scale_term() * self._mean() |
| return array_ops.matrix_set_diag( |
| -math_ops.matmul(x[..., array_ops.newaxis], |
| x[..., array_ops.newaxis, :]), # outer prod |
| self._variance()) |
| |
| def _variance(self): |
| scale = self._variance_scale_term() |
| x = scale * self._mean() |
| return x * (scale - x) |
| |
| def _variance_scale_term(self): |
| """Helper to `_covariance` and `_variance` which computes a shared scale.""" |
| return math_ops.rsqrt(1. + self.total_concentration[..., array_ops.newaxis]) |
| |
| @distribution_util.AppendDocstring( |
| """Note: The mode is undefined when any `concentration <= 1`. If |
| `self.allow_nan_stats` is `True`, `NaN` is used for undefined modes. If |
| `self.allow_nan_stats` is `False` an exception is raised when one or more |
| modes are undefined.""") |
| def _mode(self): |
| k = math_ops.cast(self.event_shape_tensor()[0], self.dtype) |
| mode = (self.concentration - 1.) / ( |
| self.total_concentration[..., array_ops.newaxis] - k) |
| if self.allow_nan_stats: |
| nan = array_ops.fill( |
| array_ops.shape(mode), |
| np.array(np.nan, dtype=self.dtype.as_numpy_dtype()), |
| name="nan") |
| return array_ops.where( |
| math_ops.reduce_all(self.concentration > 1., axis=-1), |
| mode, nan) |
| return control_flow_ops.with_dependencies([ |
| check_ops.assert_less( |
| array_ops.ones([], self.dtype), |
| self.concentration, |
| message="Mode undefined when any concentration <= 1"), |
| ], mode) |
| |
| def _maybe_assert_valid_concentration(self, concentration, validate_args): |
| """Checks the validity of the concentration parameter.""" |
| if not validate_args: |
| return concentration |
| return control_flow_ops.with_dependencies([ |
| check_ops.assert_positive( |
| concentration, |
| message="Concentration parameter must be positive."), |
| check_ops.assert_rank_at_least( |
| concentration, 1, |
| message="Concentration parameter must have >=1 dimensions."), |
| check_ops.assert_less( |
| 1, array_ops.shape(concentration)[-1], |
| message="Concentration parameter must have event_size >= 2."), |
| ], concentration) |
| |
| def _maybe_assert_valid_sample(self, x): |
| """Checks the validity of a sample.""" |
| if not self.validate_args: |
| return x |
| return control_flow_ops.with_dependencies([ |
| check_ops.assert_positive(x, message="samples must be positive"), |
| check_ops.assert_near( |
| array_ops.ones([], dtype=self.dtype), |
| math_ops.reduce_sum(x, -1), |
| message="sample last-dimension must sum to `1`"), |
| ], x) |
| |
| |
| @kullback_leibler.RegisterKL(Dirichlet, Dirichlet) |
| def _kl_dirichlet_dirichlet(d1, d2, name=None): |
| """Batchwise KL divergence KL(d1 || d2) with d1 and d2 Dirichlet. |
| |
| Args: |
| d1: instance of a Dirichlet distribution object. |
| d2: instance of a Dirichlet distribution object. |
| name: (optional) Name to use for created operations. |
| default is "kl_dirichlet_dirichlet". |
| |
| Returns: |
| Batchwise KL(d1 || d2) |
| """ |
| with ops.name_scope(name, "kl_dirichlet_dirichlet", values=[ |
| d1.concentration, d2.concentration]): |
| # The KL between Dirichlet distributions can be derived as follows. We have |
| # |
| # Dir(x; a) = 1 / B(a) * prod_i[x[i]^(a[i] - 1)] |
| # |
| # where B(a) is the multivariate Beta function: |
| # |
| # B(a) = Gamma(a[1]) * ... * Gamma(a[n]) / Gamma(a[1] + ... + a[n]) |
| # |
| # The KL is |
| # |
| # KL(Dir(x; a), Dir(x; b)) = E_Dir(x; a){log(Dir(x; a) / Dir(x; b))} |
| # |
| # so we'll need to know the log density of the Dirichlet. This is |
| # |
| # log(Dir(x; a)) = sum_i[(a[i] - 1) log(x[i])] - log B(a) |
| # |
| # The only term that matters for the expectations is the log(x[i]). To |
| # compute the expectation of this term over the Dirichlet density, we can |
| # use the following facts about the Dirichlet in exponential family form: |
| # 1. log(x[i]) is a sufficient statistic |
| # 2. expected sufficient statistics (of any exp family distribution) are |
| # equal to derivatives of the log normalizer with respect to |
| # corresponding natural parameters: E{T[i](x)} = dA/d(eta[i]) |
| # |
| # To proceed, we can rewrite the Dirichlet density in exponential family |
| # form as follows: |
| # |
| # Dir(x; a) = exp{eta(a) . T(x) - A(a)} |
| # |
| # where '.' is the dot product of vectors eta and T, and A is a scalar: |
| # |
| # eta[i](a) = a[i] - 1 |
| # T[i](x) = log(x[i]) |
| # A(a) = log B(a) |
| # |
| # Now, we can use fact (2) above to write |
| # |
| # E_Dir(x; a)[log(x[i])] |
| # = dA(a) / da[i] |
| # = d/da[i] log B(a) |
| # = d/da[i] (sum_j lgamma(a[j])) - lgamma(sum_j a[j]) |
| # = digamma(a[i])) - digamma(sum_j a[j]) |
| # |
| # Putting it all together, we have |
| # |
| # KL[Dir(x; a) || Dir(x; b)] |
| # = E_Dir(x; a){log(Dir(x; a) / Dir(x; b)} |
| # = E_Dir(x; a){sum_i[(a[i] - b[i]) log(x[i])} - (lbeta(a) - lbeta(b)) |
| # = sum_i[(a[i] - b[i]) * E_Dir(x; a){log(x[i])}] - lbeta(a) + lbeta(b) |
| # = sum_i[(a[i] - b[i]) * (digamma(a[i]) - digamma(sum_j a[j]))] |
| # - lbeta(a) + lbeta(b)) |
| |
| digamma_sum_d1 = math_ops.digamma( |
| math_ops.reduce_sum(d1.concentration, axis=-1, keepdims=True)) |
| digamma_diff = math_ops.digamma(d1.concentration) - digamma_sum_d1 |
| concentration_diff = d1.concentration - d2.concentration |
| |
| return (math_ops.reduce_sum(concentration_diff * digamma_diff, axis=-1) - |
| special_math_ops.lbeta(d1.concentration) + |
| special_math_ops.lbeta(d2.concentration)) |
