# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Random Walk Metropolis (RWM) Transition Kernel."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
# Dependency imports

import tensorflow as tf

from tensorflow_probability.python import distributions
from tensorflow_probability.python.mcmc import kernel as kernel_base
from tensorflow_probability.python.mcmc import metropolis_hastings
from tensorflow_probability.python.mcmc.internal import util as mcmc_util

__all__ = [
    'random_walk_normal_fn',
    'random_walk_uniform_fn',
    'RandomWalkMetropolis',
    'UncalibratedRandomWalk',
]

UncalibratedRandomWalkResults = collections.namedtuple(
    'UncalibratedRandomWalkResults',
    [
        'log_acceptance_correction',
        'target_log_prob',        # For "next_state".
    ])

def random_walk_normal_fn(scale=1., name=None):
  """Returns a callable that adds a random normal perturbation to the input.

  This function returns a callable that accepts a Python `list` of `Tensor`s of
  any shapes and `dtypes`  representing the state parts of the `current_state`
  and a random seed. The supplied argument `scale` must be a `Tensor` or Python
  `list` of `Tensor`s representing the scale of the generated
  proposal. `scale` must broadcast with the state parts of `current_state`.
  The callable adds a sample from a zero-mean normal distribution with the
  supplied scales to each state part and returns a same-type `list` of `Tensor`s
  as the state parts of `current_state`.

  Args:
    scale: a `Tensor` or Python `list` of `Tensor`s of any shapes and `dtypes`
      controlling the scale of the normal proposal distribution.
    name: Python `str` name prefixed to Ops created by this function.
        Default value: 'random_walk_normal_fn'.

  Returns:
    random_walk_normal_fn: A callable accepting a Python `list` of `Tensor`s
      representing the state parts of the `current_state` and an `int`
      representing the random seed to be used to generate the proposal. The
      callable returns the same-type `list` of `Tensor`s as the input and
      represents the proposal for the RWM algorithm.
  """
  def _fn(state_parts, seed):
    """Adds a normal perturbation to the input state.

    Args:
      state_parts: A list of `Tensor`s of any shape and real dtype representing
        the state parts of the `current_state` of the Markov chain.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.

    Returns:
      perturbed_state_parts: A Python `list` of The `Tensor`s. Has the same
        shape and type as the `state_parts`.

    Raises:
      ValueError: if `scale` does not broadcast with `state_parts`.
    """
    with tf.compat.v1.name_scope(
        name, 'random_walk_normal_fn', values=[state_parts, scale, seed]):
      scales = scale if mcmc_util.is_list_like(scale) else [scale]
      if len(scales) == 1:
        scales *= len(state_parts)
      if len(state_parts) != len(scales):
        raise ValueError('`scale` must broadcast with `state_parts`.')
      seed_stream = distributions.SeedStream(seed, salt='RandomWalkNormalFn')
      next_state_parts = [
          tf.random.normal(
              mean=state_part,
              stddev=scale_part,
              shape=tf.shape(input=state_part),
              dtype=state_part.dtype.base_dtype,
              seed=seed_stream())
          for scale_part, state_part in zip(scales, state_parts)
      ]

      return next_state_parts
  return _fn

def random_walk_uniform_fn(scale=1., name=None):
  """Returns a callable that adds a random uniform perturbation to the input.

  For more details on `random_walk_uniform_fn`, see
  `random_walk_normal_fn`. `scale` might
  be a `Tensor` or a list of `Tensor`s that should broadcast with state parts
  of the `current_state`. The generated uniform perturbation is sampled as a
  uniform point on the rectangle `[-scale, scale]`.

  Args:
    scale: a `Tensor` or Python `list` of `Tensor`s of any shapes and `dtypes`
      controlling the upper and lower bound of the uniform proposal
      distribution.
    name: Python `str` name prefixed to Ops created by this function.
        Default value: 'random_walk_uniform_fn'.

  Returns:
    random_walk_uniform_fn: A callable accepting a Python `list` of `Tensor`s
      representing the state parts of the `current_state` and an `int`
      representing the random seed used to generate the proposal. The callable
      returns the same-type `list` of `Tensor`s as the input and represents the
      proposal for the RWM algorithm.
  """
  def _fn(state_parts, seed):
    """Adds a uniform perturbation to the input state.

    Args:
      state_parts: A list of `Tensor`s of any shape and real dtype representing
        the state parts of the `current_state` of the Markov chain.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.

    Returns:
      perturbed_state_parts: A Python `list` of The `Tensor`s. Has the same
        shape and type as the `state_parts`.

    Raises:
      ValueError: if `scale` does not broadcast with `state_parts`.
    """
    with tf.compat.v1.name_scope(
        name, 'random_walk_uniform_fn', values=[state_parts, scale, seed]):
      scales = scale if mcmc_util.is_list_like(scale) else [scale]
      if len(scales) == 1:
        scales *= len(state_parts)
      if len(state_parts) != len(scales):
        raise ValueError('`scale` must broadcast with `state_parts`.')
      seed_stream = distributions.SeedStream(seed, salt='RandomWalkUniformFn')
      next_state_parts = [
          tf.random.uniform(
              minval=state_part - scale_part,
              maxval=state_part + scale_part,
              shape=tf.shape(input=state_part),
              dtype=state_part.dtype.base_dtype,
              seed=seed_stream())
          for scale_part, state_part in zip(scales, state_parts)
      ]
      return next_state_parts
  return _fn

class RandomWalkMetropolis(kernel_base.TransitionKernel):
  """Runs one step of the RWM algorithm with symmetric proposal.

  Random Walk Metropolis is a gradient-free Markov chain Monte Carlo
  (MCMC) algorithm. The algorithm involves a proposal generating step
  `proposal_state = current_state + perturb` by a random
  perturbation, followed by Metropolis-Hastings accept/reject step. For more
  details see [Section 2.1 of Roberts and Rosenthal (2004)](
  http://emis.ams.org/journals/PS/images/getdoc510c.pdf?id=35&article=15&mode=pdf).

  Current class implements RWM for normal and uniform proposals. Alternatively,
  the user can supply any custom proposal generating function.

  The function `one_step` can update multiple chains in parallel. It assumes
  that all leftmost dimensions of `current_state` index independent chain states
  (and are therefore updated independently). The output of
  `target_log_prob_fn(*current_state)` should sum log-probabilities across all
  event dimensions. Slices along the rightmost dimensions may have different
  target distributions; for example, `current_state[0, :]` could have a
  different target distribution from `current_state[1, :]`. These semantics
  are governed by `target_log_prob_fn(*current_state)`. (The number of
  independent chains is `tf.size(target_log_prob_fn(*current_state))`.)

  #### Examples:

  ##### Sampling from the Standard Normal Distribution.

  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  dtype = np.float32

  target = tfd.Normal(loc=dtype(0), scale=dtype(1))

  samples, _ = tfp.mcmc.sample_chain(
    num_results=1000,
    current_state=dtype(1),
    kernel=tfp.mcmc.RandomWalkMetropolis(
       target.log_prob,
       seed=42),
    num_burnin_steps=500,
    parallel_iterations=1)  # For determinism.

  sample_mean = tf.math.reduce_mean(samples, axis=0)
  sample_std = tf.sqrt(
      tf.math.reduce_mean(tf.squared_difference(samples, sample_mean),
                     axis=0))
  with tf.Session() as sess:
    [sample_mean_, sample_std_] = sess.run([sample_mean, sample_std])

  print('Estimated mean: {}'.format(sample_mean_))
  print('Estimated standard deviation: {}'.format(sample_std_))
  ```

  ##### Sampling from a 2-D Normal Distribution.

  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp

  tfd = tfp.distributions

  dtype = np.float32
  true_mean = dtype([0, 0])
  true_cov = dtype([[1, 0.5],
                   [0.5, 1]])
  num_results = 500
  num_chains = 100

  # Target distribution is defined through the Cholesky decomposition `L`:
  L = tf.linalg.cholesky(true_cov)
  target = tfd.MultivariateNormalTriL(loc=true_mean, scale_tril=L)

  # Assume that the state is passed as a list of 1-d tensors `x` and `y`.
  # Then the target log-density is defined as follows:
  def target_log_prob(x, y):
    # Stack the input tensors together
    z = tf.stack([x, y], axis=-1)
    return target.log_prob(tf.squeeze(z))

  # Initial state of the chain
  init_state = [np.ones([num_chains, 1], dtype=dtype),
                np.ones([num_chains, 1], dtype=dtype)]

  # Run Random Walk Metropolis with normal proposal for `num_results`
  # iterations for `num_chains` independent chains:
  samples, _ = tfp.mcmc.sample_chain(
      num_results=num_results,
      current_state=init_state,
      kernel=tfp.mcmc.RandomWalkMetropolis(
          target_log_prob_fn=target_log_prob,
          seed=54),
      num_burnin_steps=200,
      num_steps_between_results=1,  # Thinning.
      parallel_iterations=1)
  samples = tf.stack(samples, axis=-1)

  sample_mean = tf.math.reduce_mean(samples, axis=0)
  x = tf.squeeze(samples - sample_mean)
  sample_cov = tf.matmul(tf.transpose(x, [1, 2, 0]),
                         tf.transpose(x, [1, 0, 2])) / num_results

  mean_sample_mean = tf.math.reduce_mean(sample_mean)
  mean_sample_cov = tf.math.reduce_mean(sample_cov, axis=0)
  x = tf.reshape(sample_cov - mean_sample_cov, [num_chains, 2 * 2])
  cov_sample_cov = tf.reshape(tf.matmul(x, x, transpose_a=True) / num_chains,
                              shape=[2 * 2, 2 * 2])

  with tf.Session() as sess:
    [
      mean_sample_mean_,
      mean_sample_cov_,
      cov_sample_cov_,
    ] = sess.run([
      mean_sample_mean,
      mean_sample_cov,
      cov_sample_cov,
    ])

  print('Estimated mean: {}'.format(mean_sample_mean_))
  print('Estimated avg covariance: {}'.format(mean_sample_cov_))
  print('Estimated covariance of covariance: {}'.format(cov_sample_cov_))
  ```

  ##### Sampling from the Standard Normal Distribution using Cauchy proposal.

  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp

  tfd = tfp.distributions

  dtype = np.float32
  num_burnin_steps = 500
  num_chain_results = 1000

  def cauchy_new_state_fn(scale, dtype):
    cauchy = tfd.Cauchy(loc=dtype(0), scale=dtype(scale))
    def _fn(state_parts, seed):
      next_state_parts = []
      seed_stream  = tfd.SeedStream(seed, salt='RandomCauchy')
      for sp in state_parts:
        next_state_parts.append(sp + cauchy.sample(
          sample_shape=sp.shape, seed=seed_stream()))
      return next_state_parts
    return _fn

  target = tfd.Normal(loc=dtype(0), scale=dtype(1))

  samples, _ = tfp.mcmc.sample_chain(
      num_results=num_chain_results,
      num_burnin_steps=num_burnin_steps,
      current_state=dtype(1),
      kernel=tfp.mcmc.RandomWalkMetropolis(
          target.log_prob,
          new_state_fn=cauchy_new_state_fn(scale=0.5, dtype=dtype),
          seed=42),
      parallel_iterations=1)  # For determinism.

  sample_mean = tf.math.reduce_mean(samples, axis=0)
  sample_std = tf.sqrt(
      tf.math.reduce_mean(tf.squared_difference(samples, sample_mean),
                     axis=0))
  with tf.Session() as sess:
    [sample_mean_, sample_std_] = sess.run([sample_mean, sample_std])

  print('Estimated mean: {}'.format(sample_mean_))
  print('Estimated standard deviation: {}'.format(sample_std_))
  ```

  """

  def __init__(self,
               target_log_prob_fn,
               new_state_fn=None,
               seed=None,
               name=None):
    """Initializes this transition kernel.

    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      new_state_fn: Python callable which takes a list of state parts and a
        seed; returns a same-type `list` of `Tensor`s, each being a perturbation
        of the input state parts. The perturbation distribution is assumed to be
        a symmetric distribution centered at the input state part.
        Default value: `None` which is mapped to
          `tfp.mcmc.random_walk_normal_fn()`.
      seed: Python integer to seed the random number generator.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'rwm_kernel').

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) at each result step. Has same shape as
        `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.

    Raises:
      ValueError: if there isn't one `scale` or a list with same length as
        `current_state`.
    """
    if new_state_fn is None:
      new_state_fn = random_walk_normal_fn()

    self._impl = metropolis_hastings.MetropolisHastings(
        inner_kernel=UncalibratedRandomWalk(
            target_log_prob_fn=target_log_prob_fn,
            new_state_fn=new_state_fn,
            seed=seed,
            name=name),
        seed=seed)

  @property
  def target_log_prob_fn(self):
    return self._impl.inner_kernel.target_log_prob_fn

  @property
  def new_state_fn(self):
    return self._impl.inner_kernel.new_state_fn

  @property
  def seed(self):
    return self._impl.inner_kernel.seed

  @property
  def name(self):
    return self._impl.inner_kernel.name

  @property
  def is_calibrated(self):
    return True

  @property
  def parameters(self):
    """Return `dict` of ``__init__`` arguments and their values."""
    return self._impl.inner_kernel.parameters

  def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of Random Walk Metropolis with normal proposal.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.

    Raises:
      ValueError: if there isn't one `scale` or a list with same length as
        `current_state`.
    """
    return self._impl.one_step(current_state, previous_kernel_results)

  def bootstrap_results(self, init_state):
    """Creates initial `previous_kernel_results` using a supplied `state`."""
    return self._impl.bootstrap_results(init_state)

class UncalibratedRandomWalk(kernel_base.TransitionKernel):
  """Generate proposal for the Random Walk Metropolis algorithm.

  Warning: this kernel will not result in a chain which converges to the
  `target_log_prob`. To get a convergent MCMC, use
  `tfp.mcmc.RandomWalkMetropolisNormal(...)` or
  `tfp.mcmc.MetropolisHastings(tfp.mcmc.UncalibratedRandomWalk(...))`.

  For more details on `UncalibratedRandomWalk`, see
  `RandomWalkMetropolis`.
  """

  @mcmc_util.set_doc(RandomWalkMetropolis.__init__.__doc__)
  def __init__(self,
               target_log_prob_fn,
               new_state_fn=None,
               seed=None,
               name=None):
    if new_state_fn is None:
      new_state_fn = random_walk_normal_fn()

    self._target_log_prob_fn = target_log_prob_fn
    self._seed_stream = distributions.SeedStream(
        seed, salt='RandomWalkMetropolis')
    self._name = name
    self._parameters = dict(
        target_log_prob_fn=target_log_prob_fn,
        new_state_fn=new_state_fn,
        seed=seed,
        name=name)

  @property
  def target_log_prob_fn(self):
    return self._parameters['target_log_prob_fn']

  @property
  def new_state_fn(self):
    return self._parameters['new_state_fn']

  @property
  def seed(self):
    return self._parameters['seed']

  @property
  def name(self):
    return self._parameters['name']

  @property
  def parameters(self):
    """Return `dict` of ``__init__`` arguments and their values."""
    return self._parameters

  @property
  def is_calibrated(self):
    return False

  @mcmc_util.set_doc(RandomWalkMetropolis.one_step.__doc__)
  def one_step(self, current_state, previous_kernel_results):
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'rwm', 'one_step'),
        values=[
            self.seed, current_state, previous_kernel_results.target_log_prob
        ]):
      with tf.compat.v1.name_scope('initialize'):
        if mcmc_util.is_list_like(current_state):
          current_state_parts = list(current_state)
        else:
          current_state_parts = [current_state]
        current_state_parts = [
            tf.convert_to_tensor(value=s, name='current_state')
            for s in current_state_parts
        ]

      next_state_parts = self.new_state_fn(current_state_parts,  # pylint: disable=not-callable
                                           self._seed_stream())
      # Compute `target_log_prob` so its available to MetropolisHastings.
      next_target_log_prob = self.target_log_prob_fn(*next_state_parts)  # pylint: disable=not-callable

      def maybe_flatten(x):
        return x if mcmc_util.is_list_like(current_state) else x[0]

      return [
          maybe_flatten(next_state_parts),
          UncalibratedRandomWalkResults(
              log_acceptance_correction=tf.zeros(
                  shape=tf.shape(input=next_target_log_prob),
                  dtype=next_target_log_prob.dtype.base_dtype),
              target_log_prob=next_target_log_prob,
          ),
      ]

  @mcmc_util.set_doc(RandomWalkMetropolis.bootstrap_results.__doc__)
  def bootstrap_results(self, init_state):
    with tf.compat.v1.name_scope(self.name, 'rwm_bootstrap_results',
                                 [init_state]):
      if not mcmc_util.is_list_like(init_state):
        init_state = [init_state]
      init_state = [tf.convert_to_tensor(value=x) for x in init_state]
      init_target_log_prob = self.target_log_prob_fn(*init_state)  # pylint:disable=not-callable
      return UncalibratedRandomWalkResults(
          log_acceptance_correction=tf.zeros_like(init_target_log_prob),
          target_log_prob=init_target_log_prob)

def _maybe_call_fn(fn,
                   fn_arg_list,
                   fn_result=None,
                   description='target_log_prob'):
  """Helper which computes `fn_result` if needed."""
  if mcmc_util.is_list_like(fn_arg_list):
    fn_arg_list = list(fn_arg_list)
  else:
    fn_arg_list = [fn_arg_list]

  if fn_result is None:
    fn_result = fn(*fn_arg_list)
  if not fn_result.dtype.is_floating:
    raise TypeError('`{}` must be a `Tensor` with `float` `dtype`.'.format(
        description))
  return fn_result


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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 Independent distribution class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import distribution as distribution_lib
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import prefer_static
from tensorflow_probability.python.internal import tensorshape_util

class Independent(distribution_lib.Distribution):
  """Independent distribution from batch of distributions.

  This distribution is useful for regarding a collection of independent,
  non-identical distributions as a single random variable. For example, the
  `Independent` distribution composed of a collection of `Bernoulli`
  distributions might define a distribution over an image (where each
  `Bernoulli` is a distribution over each pixel).

  More precisely, a collection of `B` (independent) `E`-variate random variables
  (rv) `{X_1, ..., X_B}`, can be regarded as a `[B, E]`-variate random variable
  `(X_1, ..., X_B)` with probability
  `p(x_1, ..., x_B) = p_1(x_1) * ... * p_B(x_B)` where `p_b(X_b)` is the
  probability of the `b`-th rv. More generally `B, E` can be arbitrary shapes.

  Similarly, the `Independent` distribution specifies a distribution over `[B,
  E]`-shaped events. It operates by reinterpreting the rightmost batch dims as
  part of the event dimensions. The `reinterpreted_batch_ndims` parameter
  controls the number of batch dims which are absorbed as event dims;
  `reinterpreted_batch_ndims <= len(batch_shape)`.  For example, the `log_prob`
  function entails a `reduce_sum` over the rightmost `reinterpreted_batch_ndims`
  after calling the base distribution's `log_prob`.  In other words, since the
  batch dimension(s) index independent distributions, the resultant multivariate
  will have independent components.

  #### Mathematical Details

  The probability function is,

  ```none
  prob(x; reinterpreted_batch_ndims) = tf.reduce_prod(
      dist.prob(x),
      axis=-1-range(reinterpreted_batch_ndims))
  ```

  #### Examples

  ```python
  tfd = tfp.distributions

  # Make independent distribution from a 2-batch Normal.
  ind = tfd.Independent(
      distribution=tfd.Normal(loc=[-1., 1], scale=[0.1, 0.5]),
      reinterpreted_batch_ndims=1)

  # All batch dims have been "absorbed" into event dims.
  ind.batch_shape  # ==> []
  ind.event_shape  # ==> [2]

  # Make independent distribution from a 2-batch bivariate Normal.
  ind = tfd.Independent(
      distribution=tfd.MultivariateNormalDiag(
          loc=[[-1., 1], [1, -1]],
          scale_identity_multiplier=[1., 0.5]),
      reinterpreted_batch_ndims=1)

  # All batch dims have been "absorbed" into event dims.
  ind.batch_shape  # ==> []
  ind.event_shape  # ==> [2, 2]
  ```

  """

  def __init__(
      self, distribution, reinterpreted_batch_ndims=None,
      validate_args=False, name=None):
    """Construct a `Independent` distribution.

    Args:
      distribution: The base distribution instance to transform. Typically an
        instance of `Distribution`.
      reinterpreted_batch_ndims: Scalar, integer number of rightmost batch dims
        which will be regarded as event dims. When `None` all but the first
        batch axis (batch axis 0) will be transferred to event dimensions
        (analogous to `tf.layers.flatten`).
      validate_args: Python `bool`.  Whether to validate input with asserts.
        If `validate_args` is `False`, and the inputs are invalid,
        correct behavior is not guaranteed.
      name: The name for ops managed by the distribution.
        Default value: `Independent + distribution.name`.

    Raises:
      ValueError: if `reinterpreted_batch_ndims` exceeds
        `distribution.batch_ndims`
    """
    parameters = dict(locals())
    name = name or "Independent" + distribution.name
    self._distribution = distribution
    with tf.name_scope(name) as name:
      if reinterpreted_batch_ndims is None:
        reinterpreted_batch_ndims = self._get_default_reinterpreted_batch_ndims(
            distribution)
      reinterpreted_batch_ndims = tf.convert_to_tensor(
          value=reinterpreted_batch_ndims,
          dtype=tf.int32,
          name="reinterpreted_batch_ndims")
      self._reinterpreted_batch_ndims = reinterpreted_batch_ndims
      self._static_reinterpreted_batch_ndims = tf.get_static_value(
          reinterpreted_batch_ndims)
      if self._static_reinterpreted_batch_ndims is not None:
        self._reinterpreted_batch_ndims = self._static_reinterpreted_batch_ndims
      super(Independent, self).__init__(
          dtype=self._distribution.dtype,
          reparameterization_type=self._distribution.reparameterization_type,
          validate_args=validate_args,
          allow_nan_stats=self._distribution.allow_nan_stats,
          parameters=parameters,
          graph_parents=(
              [reinterpreted_batch_ndims] +
              distribution._graph_parents),  # pylint: disable=protected-access
          name=name)
      self._runtime_assertions = self._make_runtime_assertions(
          distribution, reinterpreted_batch_ndims, validate_args)

  @property
  def distribution(self):
    return self._distribution

  @property
  def reinterpreted_batch_ndims(self):
    return self._reinterpreted_batch_ndims

  def __getitem__(self, slices):
    # Because slicing is parameterization-dependent, we only implement slicing
    # for instances of Independent, not subclasses thereof.
    if type(self) is not Independent:  # pylint: disable=unidiomatic-typecheck
      return super(Independent, self).__getitem__(slices)

    if self._static_reinterpreted_batch_ndims is None:
      raise NotImplementedError(
          "Cannot slice Independent with non-static reinterpreted_batch_ndims")
    slices = (tuple(slices) if isinstance(slices, collections.Sequence)
              else (slices,))
    if Ellipsis not in slices:
      slices = slices + (Ellipsis,)
    slices = slices + (slice(None),) * int(
        self._static_reinterpreted_batch_ndims)
    return self.copy(
        distribution=self.distribution[slices],
        reinterpreted_batch_ndims=self._static_reinterpreted_batch_ndims)

  def _batch_shape_tensor(self):
    with tf.control_dependencies(self._runtime_assertions):
      batch_shape = self.distribution.batch_shape_tensor()
      batch_ndims = prefer_static.rank_from_shape(
          batch_shape, self.distribution.batch_shape)
      return batch_shape[:batch_ndims - self.reinterpreted_batch_ndims]

  def _batch_shape(self):
    batch_shape = self.distribution.batch_shape
    if (self._static_reinterpreted_batch_ndims is None or
        tensorshape_util.rank(batch_shape) is None):
      return tf.TensorShape(None)
    d = (tensorshape_util.rank(batch_shape) -
         self._static_reinterpreted_batch_ndims)
    return batch_shape[:d]

  def _event_shape_tensor(self):
    with tf.control_dependencies(self._runtime_assertions):
      batch_shape = self.distribution.batch_shape_tensor()
      batch_ndims = prefer_static.rank_from_shape(
          batch_shape, self.distribution.batch_shape)
      return prefer_static.concat([
          batch_shape[batch_ndims - self.reinterpreted_batch_ndims:],
          self.distribution.event_shape_tensor(),
      ], axis=0)

  def _event_shape(self):
    batch_shape = self.distribution.batch_shape
    if self._static_reinterpreted_batch_ndims is None:
      return tf.TensorShape(None)
    if tensorshape_util.rank(batch_shape) is not None:
      reinterpreted_batch_shape = batch_shape[
          tensorshape_util.rank(batch_shape) -
          self._static_reinterpreted_batch_ndims:]
    else:
      reinterpreted_batch_shape = tf.TensorShape(
          [None] * int(self._static_reinterpreted_batch_ndims))
    return tensorshape_util.concatenate(reinterpreted_batch_shape,
                                        self.distribution.event_shape)

  def _sample_n(self, n, seed, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self.distribution.sample(sample_shape=n, seed=seed, **kwargs)

  def _log_prob(self, x, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self._reduce(
          tf.reduce_sum, self.distribution.log_prob(x, **kwargs))

  def _log_cdf(self, x, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self._reduce(tf.reduce_sum, self.distribution.log_cdf(x, **kwargs))

  def _entropy(self, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self._reduce(tf.reduce_sum, self.distribution.entropy(**kwargs))

  def _mean(self, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self.distribution.mean(**kwargs)

  def _variance(self, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self.distribution.variance(**kwargs)

  def _stddev(self, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self.distribution.stddev(**kwargs)

  def _mode(self, **kwargs):
    with tf.control_dependencies(self._runtime_assertions):
      return self.distribution.mode(**kwargs)

  def _make_runtime_assertions(
      self, distribution, reinterpreted_batch_ndims, validate_args):
    assertions = []
    static_reinterpreted_batch_ndims = tf.get_static_value(
        reinterpreted_batch_ndims)
    batch_ndims = tensorshape_util.rank(distribution.batch_shape)
    if batch_ndims is not None and static_reinterpreted_batch_ndims is not None:
      if static_reinterpreted_batch_ndims > batch_ndims:
        raise ValueError("reinterpreted_batch_ndims({}) cannot exceed "
                         "distribution.batch_ndims({})".format(
                             static_reinterpreted_batch_ndims, batch_ndims))
    elif validate_args:
      assertions.append(
          assert_util.assert_less_equal(
              reinterpreted_batch_ndims,
              prefer_static.rank_from_shape(distribution.batch_shape_tensor,
                                            distribution.batch_shape),
              message=("reinterpreted_batch_ndims cannot exceed "
                       "distribution.batch_ndims")))
    return assertions

  def _reduce(self, op, stat):
    axis = 1 + prefer_static.range(self._reinterpreted_batch_ndims)
    return op(stat, axis=-axis)

  def _get_default_reinterpreted_batch_ndims(self, distribution):
    """Computes the default value for reinterpreted_batch_ndim __init__ arg."""
    ndims = prefer_static.rank_from_shape(
        distribution.batch_shape_tensor, distribution.batch_shape)
    return prefer_static.maximum(0, ndims - 1)

@kullback_leibler.RegisterKL(Independent, Independent)
def _kl_independent(a, b, name="kl_independent"):
  """Batched KL divergence `KL(a || b)` for Independent distributions.

  We can leverage the fact that
  ```
  KL(Independent(a) || Independent(b)) = sum(KL(a || b))
  ```
  where the sum is over the `reinterpreted_batch_ndims`.

  Args:
    a: Instance of `Independent`.
    b: Instance of `Independent`.
    name: (optional) name to use for created ops. Default "kl_independent".

  Returns:
    Batchwise `KL(a || b)`.

  Raises:
    ValueError: If the event space for `a` and `b`, or their underlying
      distributions don't match.
  """
  p = a.distribution
  q = b.distribution

  # The KL between any two (non)-batched distributions is a scalar.
  # Given that the KL between two factored distributions is the sum, i.e.
  # KL(p1(x)p2(y) || q1(x)q2(y)) = KL(p1 || q1) + KL(q1 || q2), we compute
  # KL(p || q) and do a `reduce_sum` on the reinterpreted batch dimensions.
  if (tensorshape_util.is_fully_defined(a.event_shape) and
      tensorshape_util.is_fully_defined(b.event_shape)):
    if a.event_shape == b.event_shape:
      if p.event_shape == q.event_shape:
        num_reduce_dims = (tensorshape_util.rank(a.event_shape) -
                           tensorshape_util.rank(p.event_shape))
        reduce_dims = [-i - 1 for i in range(0, num_reduce_dims)]

        return tf.reduce_sum(
            input_tensor=kullback_leibler.kl_divergence(p, q, name=name),
            axis=reduce_dims)
      else:
        raise NotImplementedError("KL between Independents with different "
                                  "event shapes not supported.")
    else:
      raise ValueError("Event shapes do not match.")
  else:
    with tf.control_dependencies(
        [
            assert_util.assert_equal(a.event_shape_tensor(),
                                     b.event_shape_tensor()),
            assert_util.assert_equal(p.event_shape_tensor(),
                                     q.event_shape_tensor())
        ]):
      num_reduce_dims = (
          prefer_static.rank_from_shape(
              a.event_shape_tensor, a.event_shape) -
          prefer_static.rank_from_shape(
              p.event_shape_tensor, a.event_shape))
      reduce_dims = prefer_static.range(-num_reduce_dims - 1, -1, 1)
      return tf.reduce_sum(
          input_tensor=kullback_leibler.kl_divergence(p, q, name=name),
          axis=reduce_dims)


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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 Mixture distribution class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import numpy as np
import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import categorical
from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import seed_stream
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

class Mixture(distribution.Distribution):
  """Mixture distribution.

  The `Mixture` object implements batched mixture distributions.
  The mixture model is defined by a `Categorical` distribution (the mixture)
  and a python list of `Distribution` objects.

  Methods supported include `log_prob`, `prob`, `mean`, `sample`, and
  `entropy_lower_bound`.

  #### Examples

  ```python
  # Create a mixture of two Gaussians:
  tfd = tfp.distributions
  mix = 0.3
  bimix_gauss = tfd.Mixture(
    cat=tfd.Categorical(probs=[mix, 1.-mix]),
    components=[
      tfd.Normal(loc=-1., scale=0.1),
      tfd.Normal(loc=+1., scale=0.5),
  ])

  # Plot the PDF.
  import matplotlib.pyplot as plt
  x = tf.linspace(-2., 3., int(1e4)).eval()
  plt.plot(x, bimix_gauss.prob(x).eval());
  ```

  """

  def __init__(self,
               cat,
               components,
               validate_args=False,
               allow_nan_stats=True,
               use_static_graph=False,
               name="Mixture"):
    """Initialize a Mixture distribution.

    A `Mixture` is defined by a `Categorical` (`cat`, representing the
    mixture probabilities) and a list of `Distribution` objects
    all having matching dtype, batch shape, event shape, and continuity
    properties (the components).

    The `num_classes` of `cat` must be possible to infer at graph construction
    time and match `len(components)`.

    Args:
      cat: A `Categorical` distribution instance, representing the probabilities
          of `distributions`.
      components: A list or tuple of `Distribution` instances.
        Each instance must have the same type, be defined on the same domain,
        and have matching `event_shape` and `batch_shape`.
      validate_args: Python `bool`, default `False`. If `True`, raise a runtime
        error if batch or event ranks are inconsistent between cat and any of
        the distributions. This is only checked if the ranks cannot be
        determined statically at graph construction time.
      allow_nan_stats: Boolean, default `True`. If `False`, raise an
       exception if a statistic (e.g. mean/mode/etc...) is undefined for any
        batch member. If `True`, batch members with valid parameters leading to
        undefined statistics will return NaN for this statistic.
      use_static_graph: Calls to `sample` will not rely on dynamic tensor
        indexing, allowing for some static graph compilation optimizations, but
        at the expense of sampling all underlying distributions in the mixture.
        (Possibly useful when running on TPUs).
        Default value: `False` (i.e., use dynamic indexing).
      name: A name for this distribution (optional).

    Raises:
      TypeError: If cat is not a `Categorical`, or `components` is not
        a list or tuple, or the elements of `components` are not
        instances of `Distribution`, or do not have matching `dtype`.
      ValueError: If `components` is an empty list or tuple, or its
        elements do not have a statically known event rank.
        If `cat.num_classes` cannot be inferred at graph creation time,
        or the constant value of `cat.num_classes` is not equal to
        `len(components)`, or all `components` and `cat` do not have
        matching static batch shapes, or all components do not
        have matching static event shapes.
    """
    parameters = dict(locals())
    if not isinstance(cat, categorical.Categorical):
      raise TypeError("cat must be a Categorical distribution, but saw: %s" %
                      cat)
    if not components:
      raise ValueError("components must be a non-empty list or tuple")
    if not isinstance(components, (list, tuple)):
      raise TypeError("components must be a list or tuple, but saw: %s" %
                      components)
    if not all(isinstance(c, distribution.Distribution) for c in components):
      raise TypeError(
          "all entries in components must be Distribution instances"
          " but saw: %s" % components)

    dtype = components[0].dtype
    if not all(d.dtype == dtype for d in components):
      raise TypeError("All components must have the same dtype, but saw "
                      "dtypes: %s" % [(d.name, d.dtype) for d in components])
    static_event_shape = components[0].event_shape
    static_batch_shape = cat.batch_shape
    for di, d in enumerate(components):
      if not tensorshape_util.is_compatible_with(static_batch_shape,
                                                 d.batch_shape):
        raise ValueError(
            "components[{}] batch shape must be compatible with cat "
            "shape and other component batch shapes".format(di))
      static_event_shape = tensorshape_util.merge_with(
          static_event_shape, d.event_shape)
      static_batch_shape = tensorshape_util.merge_with(
          static_batch_shape, d.batch_shape)
    if tensorshape_util.rank(static_event_shape) is None:
      raise ValueError(
          "Expected to know rank(event_shape) from components, but "
          "none of the components provide a static number of ndims")

    # Ensure that all batch and event ndims are consistent.
    with tf.name_scope(name) as name:
      num_components = cat.event_size
      static_num_components = tf.get_static_value(num_components)
      if static_num_components is None:
        raise ValueError(
            "Could not infer number of classes from cat and unable "
            "to compare this value to the number of components passed in.")
      # Possibly convert from numpy 0-D array.
      static_num_components = int(static_num_components)
      if static_num_components != len(components):
        raise ValueError("cat.num_classes != len(components): %d vs. %d" %
                         (static_num_components, len(components)))

      cat_batch_shape = cat.batch_shape_tensor()
      cat_batch_rank = tf.size(input=cat_batch_shape)
      if validate_args:
        batch_shapes = [d.batch_shape_tensor() for d in components]
        batch_ranks = [tf.size(input=bs) for bs in batch_shapes]
        check_message = ("components[%d] batch shape must match cat "
                         "batch shape")
        self._assertions = [
            assert_util.assert_equal(
                cat_batch_rank, batch_ranks[di], message=check_message % di)
            for di in range(len(components))
        ]
        self._assertions += [
            assert_util.assert_equal(
                cat_batch_shape, batch_shapes[di], message=check_message % di)
            for di in range(len(components))
        ]
      else:
        self._assertions = []

      self._cat = cat
      self._components = list(components)
      self._num_components = static_num_components
      self._static_event_shape = static_event_shape
      self._static_batch_shape = static_batch_shape

      self._use_static_graph = use_static_graph
      if use_static_graph and static_num_components is None:
        raise ValueError("Number of categories must be known statically when "
                         "`static_sample=True`.")
    # We let the Mixture distribution access _graph_parents since its arguably
    # more like a baseclass.
    graph_parents = self._cat._graph_parents  # pylint: disable=protected-access
    for c in self._components:
      graph_parents += c._graph_parents  # pylint: disable=protected-access

    super(Mixture, self).__init__(
        dtype=dtype,
        reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=graph_parents,
        name=name)

  @property
  def cat(self):
    return self._cat

  @property
  def components(self):
    return self._components

  @property
  def num_components(self):
    return self._num_components

  def _batch_shape_tensor(self):
    return self._cat.batch_shape_tensor()

  def _batch_shape(self):
    return self._static_batch_shape

  def _event_shape_tensor(self):
    return self._components[0].event_shape_tensor()

  def _event_shape(self):
    return self._static_event_shape

  def _expand_to_event_rank(self, x):
    """Expand the rank of x up to static_event_rank times for broadcasting.

    The static event rank was checked to not be None at construction time.

    Args:
      x: A tensor to expand.
    Returns:
      The expanded tensor.
    """
    expanded_x = x
    for _ in range(tensorshape_util.rank(self.event_shape)):
      expanded_x = tf.expand_dims(expanded_x, -1)
    return expanded_x

  def _mean(self):
    with tf.control_dependencies(self._assertions):
      distribution_means = [d.mean() for d in self.components]
      cat_probs = self._cat_probs(log_probs=False)
      cat_probs = [self._expand_to_event_rank(c_p) for c_p in cat_probs]
      partial_means = [
          c_p * m for (c_p, m) in zip(cat_probs, distribution_means)
      ]
      # These should all be the same shape by virtue of matching
      # batch_shape and event_shape.
      return tf.add_n(partial_means)

  def _stddev(self):
    with tf.control_dependencies(self._assertions):
      distribution_means = [d.mean() for d in self.components]
      distribution_devs = [d.stddev() for d in self.components]
      cat_probs = self._cat_probs(log_probs=False)

      stacked_means = tf.stack(distribution_means, axis=-1)
      stacked_devs = tf.stack(distribution_devs, axis=-1)
      cat_probs = [self._expand_to_event_rank(c_p) for c_p in cat_probs]
      broadcasted_cat_probs = (
          tf.stack(cat_probs, axis=-1) * tf.ones_like(stacked_means))

      batched_dev = distribution_util.mixture_stddev(
          tf.reshape(broadcasted_cat_probs, [-1, len(self.components)]),
          tf.reshape(stacked_means, [-1, len(self.components)]),
          tf.reshape(stacked_devs, [-1, len(self.components)]))

      # I.e. re-shape to list(batch_shape) + list(event_shape).
      return tf.reshape(batched_dev, tf.shape(input=broadcasted_cat_probs)[:-1])

  def _log_prob(self, x):
    with tf.control_dependencies(self._assertions):
      x = tf.convert_to_tensor(value=x, name="x")
      distribution_log_probs = [d.log_prob(x) for d in self.components]
      cat_log_probs = self._cat_probs(log_probs=True)
      final_log_probs = [
          cat_lp + d_lp
          for (cat_lp, d_lp) in zip(cat_log_probs, distribution_log_probs)
      ]
      concat_log_probs = tf.stack(final_log_probs, 0)
      log_sum_exp = tf.reduce_logsumexp(input_tensor=concat_log_probs, axis=[0])
      return log_sum_exp

  def _log_cdf(self, x):
    with tf.control_dependencies(self._assertions):
      x = tf.convert_to_tensor(value=x, name="x")
      distribution_log_cdfs = [d.log_cdf(x) for d in self.components]
      cat_log_probs = self._cat_probs(log_probs=True)
      final_log_cdfs = [
          cat_lp + d_lcdf
          for (cat_lp, d_lcdf) in zip(cat_log_probs, distribution_log_cdfs)
      ]
      concatted_log_cdfs = tf.stack(final_log_cdfs, axis=0)
      mixture_log_cdf = tf.reduce_logsumexp(
          input_tensor=concatted_log_cdfs, axis=[0])
      return mixture_log_cdf

  def _sample_n(self, n, seed=None):
    if self._use_static_graph:
      with tf.control_dependencies(self._assertions):
        # This sampling approach is almost the same as the approach used by
        # `MixtureSameFamily`. The differences are due to having a list of
        # `Distribution` objects rather than a single object, and maintaining
        # random seed management that is consistent with the non-static code
        # path.
        samples = []
        cat_samples = self.cat.sample(n, seed=seed)
        stream = seed_stream.SeedStream(seed, salt="Mixture")

        for c in range(self.num_components):
          samples.append(self.components[c].sample(n, seed=stream()))
        stack_axis = -1 - tensorshape_util.rank(self._static_event_shape)
        x = tf.stack(samples, axis=stack_axis)  # [n, B, k, E]
        npdt = dtype_util.as_numpy_dtype(x.dtype)
        mask = tf.one_hot(
            indices=cat_samples,  # [n, B]
            depth=self._num_components,  # == k
            on_value=npdt(1),
            off_value=npdt(0))  # [n, B, k]
        mask = distribution_util.pad_mixture_dimensions(
            mask, self, self._cat,
            tensorshape_util.rank(self._static_event_shape))  # [n, B, k, [1]*e]
        return tf.reduce_sum(
            input_tensor=x * mask,
            axis=stack_axis)  # [n, B, E]

    with tf.control_dependencies(self._assertions):
      n = tf.convert_to_tensor(value=n, name="n")
      static_n = tf.get_static_value(n)
      n = int(static_n) if static_n is not None else n
      cat_samples = self.cat.sample(n, seed=seed)

      static_samples_shape = cat_samples.shape
      if tensorshape_util.is_fully_defined(static_samples_shape):
        samples_shape = tensorshape_util.as_list(static_samples_shape)
        samples_size = tensorshape_util.num_elements(static_samples_shape)
      else:
        samples_shape = tf.shape(input=cat_samples)
        samples_size = tf.size(input=cat_samples)
      static_batch_shape = self.batch_shape
      if tensorshape_util.is_fully_defined(static_batch_shape):
        batch_shape = tensorshape_util.as_list(static_batch_shape)
        batch_size = tensorshape_util.num_elements(static_batch_shape)
      else:
        batch_shape = self.batch_shape_tensor()
        batch_size = tf.reduce_prod(input_tensor=batch_shape)
      static_event_shape = self.event_shape
      if tensorshape_util.is_fully_defined(static_event_shape):
        event_shape = np.array(
            tensorshape_util.as_list(static_event_shape), dtype=np.int32)
      else:
        event_shape = self.event_shape_tensor()

      # Get indices into the raw cat sampling tensor. We will
      # need these to stitch sample values back out after sampling
      # within the component partitions.
      samples_raw_indices = tf.reshape(tf.range(0, samples_size), samples_shape)

      # Partition the raw indices so that we can use
      # dynamic_stitch later to reconstruct the samples from the
      # known partitions.
      partitioned_samples_indices = tf.dynamic_partition(
          data=samples_raw_indices,
          partitions=cat_samples,
          num_partitions=self.num_components)

      # Copy the batch indices n times, as we will need to know
      # these to pull out the appropriate rows within the
      # component partitions.
      batch_raw_indices = tf.reshape(
          tf.tile(tf.range(0, batch_size), [n]), samples_shape)

      # Explanation of the dynamic partitioning below:
      #   batch indices are i.e., [0, 1, 0, 1, 0, 1]
      # Suppose partitions are:
      #     [1 1 0 0 1 1]
      # After partitioning, batch indices are cut as:
      #     [batch_indices[x] for x in 2, 3]
      #     [batch_indices[x] for x in 0, 1, 4, 5]
      # i.e.
      #     [1 1] and [0 0 0 0]
      # Now we sample n=2 from part 0 and n=4 from part 1.
      # For part 0 we want samples from batch entries 1, 1 (samples 0, 1),
      # and for part 1 we want samples from batch entries 0, 0, 0, 0
      #   (samples 0, 1, 2, 3).
      partitioned_batch_indices = tf.dynamic_partition(
          data=batch_raw_indices,
          partitions=cat_samples,
          num_partitions=self.num_components)
      samples_class = [None for _ in range(self.num_components)]

      stream = seed_stream.SeedStream(seed, salt="Mixture")

      for c in range(self.num_components):
        n_class = tf.size(input=partitioned_samples_indices[c])
        samples_class_c = self.components[c].sample(
            n_class, seed=stream())

        # Pull out the correct batch entries from each index.
        # To do this, we may have to flatten the batch shape.

        # For sample s, batch element b of component c, we get the
        # partitioned batch indices from
        # partitioned_batch_indices[c]; and shift each element by
        # the sample index. The final lookup can be thought of as
        # a matrix gather along locations (s, b) in
        # samples_class_c where the n_class rows correspond to
        # samples within this component and the batch_size columns
        # correspond to batch elements within the component.
        #
        # Thus the lookup index is
        #   lookup[c, i] = batch_size * s[i] + b[c, i]
        # for i = 0 ... n_class[c] - 1.
        lookup_partitioned_batch_indices = (
            batch_size * tf.range(n_class) + partitioned_batch_indices[c])
        samples_class_c = tf.reshape(
            samples_class_c, tf.concat([[n_class * batch_size], event_shape],
                                       0))
        samples_class_c = tf.gather(
            samples_class_c,
            lookup_partitioned_batch_indices,
            name="samples_class_c_gather")
        samples_class[c] = samples_class_c

      # Stitch back together the samples across the components.
      lhs_flat_ret = tf.dynamic_stitch(
          indices=partitioned_samples_indices, data=samples_class)
      # Reshape back to proper sample, batch, and event shape.
      ret = tf.reshape(
          lhs_flat_ret, tf.concat(
              [samples_shape, self.event_shape_tensor()], 0))
      tensorshape_util.set_shape(
          ret,
          tensorshape_util.concatenate(static_samples_shape, self.event_shape))
      return ret

  def entropy_lower_bound(self, name="entropy_lower_bound"):
    r"""A lower bound on the entropy of this mixture model.

    The bound below is not always very tight, and its usefulness depends
    on the mixture probabilities and the components in use.

    A lower bound is useful for ELBO when the `Mixture` is the variational
    distribution:

    \\(
    \log p(x) >= ELBO = \int q(z) \log p(x, z) dz + H[q]
    \\)

    where \\( p \\) is the prior distribution, \\( q \\) is the variational,
    and \\( H[q] \\) is the entropy of \\( q \\). If there is a lower bound
    \\( G[q] \\) such that \\( H[q] \geq G[q] \\) then it can be used in
    place of \\( H[q] \\).

    For a mixture of distributions \\( q(Z) = \sum_i c_i q_i(Z) \\) with
    \\( \sum_i c_i = 1 \\), by the concavity of \\( f(x) = -x \log x \\), a
    simple lower bound is:

    \\(
    \begin{align}
    H[q] & = - \int q(z) \log q(z) dz \\\
       & = - \int (\sum_i c_i q_i(z)) \log(\sum_i c_i q_i(z)) dz \\\
       & \geq - \sum_i c_i \int q_i(z) \log q_i(z) dz \\\
       & = \sum_i c_i H[q_i]
    \end{align}
    \\)

    This is the term we calculate below for \\( G[q] \\).

    Args:
      name: A name for this operation (optional).

    Returns:
      A lower bound on the Mixture's entropy.
    """
    with self._name_scope(name):
      with tf.control_dependencies(self._assertions):
        distribution_entropies = [d.entropy() for d in self.components]
        cat_probs = self._cat_probs(log_probs=False)
        partial_entropies = [
            c_p * m for (c_p, m) in zip(cat_probs, distribution_entropies)
        ]
        # These are all the same shape by virtue of matching batch_shape
        return tf.add_n(partial_entropies)

  def _cat_probs(self, log_probs):
    """Get a list of num_components batchwise probabilities."""
    which_softmax = tf.nn.log_softmax if log_probs else tf.nn.softmax
    cat_probs = which_softmax(self.cat.logits)
    cat_probs = tf.unstack(cat_probs, num=self.num_components, axis=-1)
    return cat_probs


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""FiniteDiscrete distribution class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf
from tensorflow_probability.python.distributions import categorical
from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.internal import distribution_util as dist_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization

__all__ = [
    'FiniteDiscrete',
]

class FiniteDiscrete(distribution.Distribution):
  """The finite discrete distribution.

  The FiniteDiscrete distribution is parameterized by either probabilities or
  log-probabilities of a set of `K` possible outcomes, which is defined by
  a strictly ascending list of `K` values.

  #### Mathematical Details

  The probability mass function (pmf) is,

  ```none
  pmf(x; pi, qi) = prod_j pi_j**[x == qi_j]
  ```

  #### Examples

  ```python

  # Initialize a discrete distribution with 4 possible outcomes and the 2nd
  # outcome being most likely.
  dist = FiniteDiscrete([1., 2., 4., 8.], probs=[0.1, 0.4, 0.3, 0.2])
  dist.prob(2.)
  # ==> 0.4

  # Using logits to initialize a discrete distribution with 4 possible outcomes
  # and the 2nd outcome being most likely.
  dist = FiniteDiscrete([1., 2., 4., 8.], logits=np.log([0.1, 0.4, 0.3, 0.2]))
  dist.prob(2.)
  # ==> 0.4
  ```

  """

  def __init__(self,
               outcomes,
               logits=None,
               probs=None,
               rtol=None,
               atol=None,
               validate_args=False,
               allow_nan_stats=True,
               name='FiniteDiscrete'):
    """Construct a finite discrete contribution.

    Args:
      outcomes: A 1-D floating or integer `Tensor`, representing a list of
        possible outcomes in strictly ascending order.
      logits: A floating N-D `Tensor`, `N >= 1`, representing the log
        probabilities of a set of FiniteDiscrete distributions. The first `N -
        1` dimensions index into a batch of independent distributions and the
        last dimension represents a vector of logits for each discrete value.
        Only one of `logits` or `probs` should be passed in.
      probs: A floating  N-D `Tensor`, `N >= 1`, representing the probabilities
        of a set of FiniteDiscrete distributions. The first `N - 1` dimensions
        index into a batch of independent distributions and the last dimension
        represents a vector of probabilities for each discrete value. Only one
        of `logits` or `probs` should be passed in.
      rtol: `Tensor` with same `dtype` as `outcomes`. The relative tolerance for
        floating number comparison. Only effective when `outcomes` is a floating
        `Tensor`. Default is `10 * eps`.
      atol: `Tensor` with same `dtype` as `outcomes`. The absolute tolerance for
        floating number comparison. Only effective when `outcomes` is a floating
        `Tensor`. Default is `10 * eps`.
      validate_args:  Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may 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 tf.compat.v1.name_scope(
        name, values=[outcomes, logits, probs]) as name:
      self._outcomes = tf.convert_to_tensor(value=outcomes, name='outcomes')
      if validate_args:
        assertions = _maybe_validate_args(self._outcomes, logits, probs,
                                          validate_args)
        with tf.control_dependencies(assertions):
          self._outcomes = tf.identity(self._outcomes)
      if dtype_util.is_floating(self._outcomes.dtype):
        eps = np.finfo(dtype_util.as_numpy_dtype(self._outcomes.dtype)).eps
        self._rtol = 10 * eps if rtol is None else rtol
        self._atol = 10 * eps if atol is None else atol
      else:
        self._rtol = None
        self._atol = None
      self._categorical = categorical.Categorical(
          logits=logits,
          probs=probs,
          dtype=tf.int32,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats)
    super(FiniteDiscrete, self).__init__(
        dtype=self._outcomes.dtype,
        reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[
            self._outcomes, self._categorical.logits, self._categorical.probs
        ],
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    # outcomes is currently not sliceable.
    return dict(logits=1, probs=1)

  @property
  def outcomes(self):
    return self._outcomes

  @property
  def logits(self):
    return self._categorical.logits

  @property
  def probs(self):
    return self._categorical.probs

  def _batch_shape_tensor(self):
    return self._categorical.batch_shape_tensor()

  def _batch_shape(self):
    return self._categorical.batch_shape

  def _event_shape_tensor(self):
    return tf.constant([], dtype=tf.int32)

  def _event_shape(self):
    return tf.TensorShape([])

  def _cdf(self, x):
    x = tf.convert_to_tensor(value=x, name='x')
    flat_x = tf.reshape(x, shape=[-1])
    upper_bound = tf.searchsorted(self.outcomes, values=flat_x, side='right')
    values_at_ub = tf.gather(
        self.outcomes,
        indices=tf.minimum(upper_bound,
                           dist_util.prefer_static_shape(self.outcomes)[-1] -
                           1))
    should_use_upper_bound = self._is_equal_or_close(flat_x, values_at_ub)
    indices = tf.where(should_use_upper_bound, upper_bound, upper_bound - 1)
    return self._categorical.cdf(
        tf.reshape(indices, shape=dist_util.prefer_static_shape(x)))

  def _entropy(self):
    return self._categorical.entropy()

  def _is_equal_or_close(self, a, b):
    if dtype_util.is_integer(self.outcomes.dtype):
      return tf.equal(a, b)
    return tf.abs(a - b) < self._atol + self._rtol * tf.abs(b)

  def _log_prob(self, x):
    x = tf.convert_to_tensor(value=x, name='x')
    right_indices = tf.minimum(
        tf.size(input=self.outcomes) - 1,
        tf.reshape(
            tf.searchsorted(
                self.outcomes, values=tf.reshape(x, shape=[-1]), side='right'),
            dist_util.prefer_static_shape(x)))
    use_right_indices = self._is_equal_or_close(
        x, tf.gather(self.outcomes, indices=right_indices))
    left_indices = tf.maximum(0, right_indices - 1)
    use_left_indices = self._is_equal_or_close(
        x, tf.gather(self.outcomes, indices=left_indices))
    log_probs = self._categorical.log_prob(
        tf.where(use_left_indices, left_indices, right_indices))
    should_be_neg_inf = tf.broadcast_to(
        tf.logical_not(use_left_indices | use_right_indices),
        shape=dist_util.prefer_static_shape(log_probs))
    return tf.where(
        should_be_neg_inf,
        tf.fill(
            dist_util.prefer_static_shape(should_be_neg_inf),
            dtype_util.as_numpy_dtype(log_probs.dtype)(-np.inf)), log_probs)

  def _mean(self):
    probs = self.probs
    outcomes = self.outcomes
    if dtype_util.is_integer(outcomes.dtype):
      if self._validate_args:
        outcomes = dist_util.embed_check_integer_casting_closed(
            outcomes, target_dtype=probs.dtype)
      outcomes = tf.cast(outcomes, dtype=probs.dtype)
    return tf.tensordot(outcomes, probs, axes=[[0], [-1]])

  def _mode(self):
    return tf.gather(self.outcomes, indices=self._categorical.mode())

  def _sample_n(self, n, seed=None, **distribution_kwargs):
    return tf.gather(
        self.outcomes,
        indices=self._categorical.sample(
            sample_shape=[n], seed=seed, **distribution_kwargs))

  def _variance(self):
    probs = self._categorical.probs
    outcomes = tf.broadcast_to(
        self.outcomes, shape=dist_util.prefer_static_shape(probs))
    if dtype_util.is_integer(outcomes.dtype):
      if self._validate_args:
        outcomes = dist_util.embed_check_integer_casting_closed(
            outcomes, target_dtype=probs.dtype)
      outcomes = tf.cast(outcomes, dtype=probs.dtype)
    square_d = tf.math.squared_difference(outcomes,
                                          tf.expand_dims(self.mean(), axis=-1))
    return tf.reduce_sum(input_tensor=probs * square_d, axis=-1)

def _maybe_validate_args(outcomes, logits, probs, validate_args):
  """Validate `outcomes`, `logits` and `probs`'s shapes."""
  assertions = []

  def validate_equal_last_dim(tensor_a, tensor_b, message):
    if tensor_a.shape.is_fully_defined() and tensor_b.shape.is_fully_defined():
      if tensor_a.shape[-1] != tensor_b.shape[-1]:
        raise ValueError(message)
    elif validate_args:
      assertions.append(
          tf.compat.v1.assert_equal(
              tf.shape(input=tensor_a)[-1],
              tf.shape(input=tensor_b)[-1],
              message=message))

  if logits is not None:
    validate_equal_last_dim(
        outcomes,
        logits,
        message='Last dimension of outcomes and logits must be equal size.')
  if probs is not None:
    validate_equal_last_dim(
        outcomes,
        probs,
        message='Last dimension of outcomes and probs must be equal size.')

  message = 'Rank of outcomes must be 1.'
  if outcomes.shape.ndims is not None:
    if outcomes.shape.ndims != 1:
      raise ValueError(message)
  elif validate_args:
    assertions.append(tf.compat.v1.assert_rank(outcomes, 1, message=message))

  message = 'Size of outcomes must be greater than 0.'
  if outcomes.shape.num_elements() is not None:
    if outcomes.shape.num_elements() == 0:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_greater(
            tf.size(input=outcomes), 0, message=message))

  if validate_args:
    assertions.append(
        tf.compat.v1.assert_equal(
            tf.math.is_strictly_increasing(outcomes),
            True,
            message='outcomes is not strictly increasing.'))

  return assertions


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Tools for probabilistic reasoning in TensorFlow."""

# Contributors to the `python/` dir should not alter this file; instead update
# `python/__init__.py` as necessary.

# We need to put some imports inside a function call below, and the function
# call needs to come before the *actual* imports that populate the
# tensorflow_probability namespace. Hence, we disable this lint check throughout
# the file.
#
# pylint: disable=g-import-not-at-top

# Ensure TensorFlow is importable and its version is sufficiently recent. This
# needs to happen before anything else, since the imports below will try to
# import tensorflow, too.
def _ensure_tf_install():  # pylint: disable=g-statement-before-imports
  """Attempt to import tensorflow, and ensure its version is sufficient.

  Raises:
    ImportError: if either tensorflow is not importable or its version is
    inadequate.
  """
  try:
    import tensorflow as tf
  except ImportError:
    # Print more informative error message, then reraise.
    print("\n\nFailed to import TensorFlow. Please note that TensorFlow is not "
          "installed by default when you install TensorFlow Probability. This "
          "is so that users can decide whether to install the GPU-enabled "
          "TensorFlow package. To use TensorFlow Probability, please install "
          "the most recent version of TensorFlow, by following instructions at "
          "https://tensorflow.org/install.\n\n")
    raise

  import distutils.version

  #
  # Update this whenever we need to depend on a newer TensorFlow release.
  #
  required_tensorflow_version = "1.13"

  if (distutils.version.LooseVersion(tf.__version__) <
      distutils.version.LooseVersion(required_tensorflow_version)):
    raise ImportError(
        "This version of TensorFlow Probability requires TensorFlow "
        "version >= {required}; Detected an installation of version {present}. "
        "Please upgrade TensorFlow to proceed.".format(
            required=required_tensorflow_version,
            present=tf.__version__))

_ensure_tf_install()

# Cleanup symbols to avoid polluting namespace.
import sys as _sys
for symbol in ["_ensure_tf_install", "_sys"]:
  delattr(_sys.modules[__name__], symbol)

# from tensorflow_probability.google import staging  # DisableOnExport
from tensorflow_probability.python import *  # pylint: disable=wildcard-import
from tensorflow_probability.python.version import __version__
# pylint: enable=g-import-not-at-top


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Trains a Bayesian logistic regression model using the No-U-Turn Sampler."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
# Dependency imports
from absl import flags
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
import numpy as np
import tensorflow as tf

from tensorflow_probability import edward2 as ed
from experimental import no_u_turn_sampler

flags.DEFINE_integer(
    "max_steps",
    default=5,
    help="Number of training steps to run.")

flags.DEFINE_string(
    "model_dir",
    default=os.path.join(os.getenv("TEST_TMPDIR", "/tmp"),
                         "logistic_regression/"),
    help="Path to write plots to.")

flags.DEFINE_bool(
    "fake_data",
    default=False,
    help="If True, uses fake data. Defaults to real data.")

flags.DEFINE_bool(
    "skip_plots",
    default=False,
    help=("If True, skips drawing plots (just checking "
          "error-free execution). Defaults to drawing plots."))

FLAGS = flags.FLAGS

def logistic_regression(features):
  """Bayesian logistic regression, which returns labels given features."""
  coeffs = ed.MultivariateNormalDiag(
      loc=tf.zeros(features.shape[1]), name="coeffs")
  labels = ed.Bernoulli(
      logits=tf.tensordot(features, coeffs, [[1], [0]]), name="labels")
  return labels

def covertype():
  """Builds the Covertype data set."""
  import sklearn.datasets  # pylint: disable=g-import-not-at-top
  data = sklearn.datasets.covtype.fetch_covtype()
  features = data.data
  labels = data.target

  # Normalize features and append a column of ones for the intercept.
  features -= features.mean(0)
  features /= features.std(0)
  features = np.hstack([features, np.ones([features.shape[0], 1])])
  features = tf.cast(features, dtype=tf.float32)

  # Binarize outcomes on whether it is a specific category.
  _, counts = np.unique(labels, return_counts=True)
  specific_category = np.argmax(counts)
  labels = (labels == specific_category)
  labels = tf.cast(labels, dtype=tf.int32)
  return features, labels

def main(argv):
  del argv  # unused
  if not FLAGS.skip_plots:
    if tf.io.gfile.exists(FLAGS.model_dir):
      tf.compat.v1.logging.warning(
          "Warning: deleting old log directory at {}".format(FLAGS.model_dir))
      tf.io.gfile.rmtree(FLAGS.model_dir)
    tf.io.gfile.makedirs(FLAGS.model_dir)

  tf.compat.v1.enable_eager_execution()
  print("Number of available GPUs", tf.contrib.eager.num_gpus())

  if FLAGS.fake_data:
    features = tf.random.normal([20, 55])
    labels = tf.random.uniform([20], minval=0, maxval=2, dtype=tf.int32)
  else:
    features, labels = covertype()
  print("Data set size", features.shape[0])
  print("Number of features", features.shape[1])

  log_joint = ed.make_log_joint_fn(logistic_regression)
  @tf.function
  def target_log_prob_fn(coeffs):
    return log_joint(features=features, coeffs=coeffs, labels=labels)

  # Initialize using a sample from 20 steps of NUTS. This is roughly a MAP
  # estimate and is written explicitly to avoid differences in warm-starts
  # between different implementations (e.g., Stan, PyMC3).
  coeffs = tf.constant(
      [+2.03420663e+00, -3.53567265e-02, -1.49223924e-01, -3.07049364e-01,
       -1.00028366e-01, -1.46827862e-01, -1.64167881e-01, -4.20344204e-01,
       +9.47479829e-02, -1.12681836e-02, +2.64442056e-01, -1.22087866e-01,
       -6.00568838e-02, -3.79419506e-01, -1.06668741e-01, -2.97053963e-01,
       -2.05253899e-01, -4.69537191e-02, -2.78072730e-02, -1.43250525e-01,
       -6.77954629e-02, -4.34899796e-03, +5.90927452e-02, +7.23133609e-02,
       +1.38526391e-02, -1.24497898e-01, -1.50733739e-02, -2.68872194e-02,
       -1.80925727e-02, +3.47936489e-02, +4.03552800e-02, -9.98773426e-03,
       +6.20188080e-02, +1.15002751e-01, +1.32145107e-01, +2.69109547e-01,
       +2.45785132e-01, +1.19035013e-01, -2.59744357e-02, +9.94279515e-04,
       +3.39266285e-02, -1.44057125e-02, -6.95222765e-02, -7.52013028e-02,
       +1.21171586e-01, +2.29205526e-02, +1.47308692e-01, -8.34354162e-02,
       -9.34122875e-02, -2.97472421e-02, -3.03937674e-01, -1.70958012e-01,
       -1.59496680e-01, -1.88516974e-01, -1.20889175e+00])

  # Initialize step size via result of 50 warmup steps from Stan.
  step_size = 0.00167132

  coeffs_samples = []
  target_log_prob = None
  grads_target_log_prob = None
  for step in range(FLAGS.max_steps):
    print("Step", step)
    [
        [coeffs],
        target_log_prob,
        grads_target_log_prob,
    ] = no_u_turn_sampler.kernel(
        target_log_prob_fn=target_log_prob_fn,
        current_state=[coeffs],
        step_size=[step_size],
        seed=step,
        current_target_log_prob=target_log_prob,
        current_grads_target_log_prob=grads_target_log_prob)
    coeffs_samples.append(coeffs)

  if not FLAGS.skip_plots:
    for coeffs_sample in coeffs_samples:
      plt.plot(coeffs_sample.numpy())

    filename = os.path.join(FLAGS.model_dir, "coeffs_samples.png")
    plt.savefig(filename)
    print("Figure saved as", filename)
    plt.close()

if __name__ == "__main__":
  tf.compat.v1.app.run()


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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

# Dependency imports
import numpy as np
import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

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

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
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  # 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 = tfd.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 = tfd.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 = tfd.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 tf.name_scope(name) as name:
      self._concentration = self._maybe_assert_valid_concentration(
          tf.convert_to_tensor(
              value=concentration,
              name="concentration",
              dtype=dtype_util.common_dtype([concentration],
                                            preferred_dtype=tf.float32)),
          validate_args)
      self._total_concentration = tf.reduce_sum(
          input_tensor=self._concentration, axis=-1)
    super(Dirichlet, self).__init__(
        dtype=self._concentration.dtype,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        parameters=parameters,
        graph_parents=[self._concentration,
                       self._total_concentration],
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(concentration=1)

  @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 tf.shape(input=self.total_concentration)

  def _batch_shape(self):
    return self.total_concentration.shape

  def _event_shape_tensor(self):
    return tf.shape(input=self.concentration)[-1:]

  def _event_shape(self):
    return tensorshape_util.with_rank_at_least(self.concentration.shape, 1)[-1:]

  def _sample_n(self, n, seed=None):
    gamma_sample = tf.random.gamma(
        shape=[n], alpha=self.concentration, dtype=self.dtype, seed=seed)
    return gamma_sample / tf.reduce_sum(
        input_tensor=gamma_sample, axis=-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 tf.exp(self._log_prob(x))

  def _log_unnormalized_prob(self, x):
    x = self._maybe_assert_valid_sample(x)
    return tf.reduce_sum(
        input_tensor=tf.math.xlogy(self.concentration - 1., x), axis=-1)

  def _log_normalization(self):
    return tf.math.lbeta(self.concentration)

  def _entropy(self):
    k = tf.cast(self.event_shape_tensor()[0], self.dtype)
    return (self._log_normalization() +
            ((self.total_concentration - k) *
             tf.math.digamma(self.total_concentration)) - tf.reduce_sum(
                 input_tensor=(self.concentration - 1.) *
                 tf.math.digamma(self.concentration),
                 axis=-1))

  def _mean(self):
    return self.concentration / self.total_concentration[..., tf.newaxis]

  def _covariance(self):
    x = self._variance_scale_term() * self._mean()
    return tf.linalg.set_diag(
        -tf.matmul(x[..., tf.newaxis], x[..., tf.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 tf.math.rsqrt(1. + self.total_concentration[..., tf.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 = tf.cast(self.event_shape_tensor()[0], self.dtype)
    mode = (self.concentration - 1.) / (
        self.total_concentration[..., tf.newaxis] - k)
    if self.allow_nan_stats:
      nan = tf.fill(
          tf.shape(input=mode),
          dtype_util.as_numpy_dtype(self.dtype)(np.nan),
          name="nan")
      return tf.where(
          tf.reduce_all(input_tensor=self.concentration > 1., axis=-1), mode,
          nan)
    return distribution_util.with_dependencies([
        assert_util.assert_less(
            tf.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 distribution_util.with_dependencies([
        assert_util.assert_positive(
            concentration, message="Concentration parameter must be positive."),
        assert_util.assert_rank_at_least(
            concentration,
            1,
            message="Concentration parameter must have >=1 dimensions."),
        assert_util.assert_less(
            1,
            tf.shape(input=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 distribution_util.with_dependencies([
        assert_util.assert_positive(x, message="samples must be positive"),
        assert_util.assert_near(
            tf.ones([], dtype=self.dtype),
            tf.reduce_sum(input_tensor=x, axis=-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 tf.name_scope(name or "kl_dirichlet_dirichlet"):
    # 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 = tf.math.digamma(
        tf.reduce_sum(input_tensor=d1.concentration, axis=-1, keepdims=True))
    digamma_diff = tf.math.digamma(d1.concentration) - digamma_sum_d1
    concentration_diff = d1.concentration - d2.concentration

    return (
        tf.reduce_sum(input_tensor=concentration_diff * digamma_diff, axis=-1) -
        tf.math.lbeta(d1.concentration) + tf.math.lbeta(d2.concentration))