# 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.
# ============================================================================
"""Blockwise bijector."""

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

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.bijectors import bijector as bijector_base
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    'Blockwise',
]

class Blockwise(bijector_base.Bijector):
  """Bijector which applies a list of bijectors to blocks of a `Tensor`.

  More specifically, given [F_0, F_1, ... F_n] which are scalar or vector
  bijectors this bijector creates a transformation which operates on the vector
  [x_0, ... x_n] with the transformation [F_0(x_0), F_1(x_1) ..., F_n(x_n)]
  where x_0, ..., x_n are blocks (partitions) of the vector.

  Example Use:

  ```python
  blockwise = tfb.Blockwise(
      bijectors=[tfb.Exp(), tfb.Sigmoid()], block_sizes=[2, 1]
    )
  y = blockwise.forward(x)

  # Equivalent to:
  x_0, x_1 = tf.split(x, [2, 1], axis=-1)
  y_0 = tfb.Exp().forward(x_0)
  y_1 = tfb.Sigmoid().forward(x_1)
  y = tf.concat([y_0, y_1], axis=-1)
  ```
  """

  def __init__(self,
               bijectors,
               block_sizes=None,
               validate_args=False,
               name=None):
    """Creates the bijector.

    Args:
      bijectors: A non-empty list of bijectors.
      block_sizes: A 1-D integer `Tensor` with each element signifying the
        length of the block of the input vector to pass to the corresponding
        bijector. The length of `block_sizes` must be be equal to the length of
        `bijectors`. If left as None, a vector of 1's is used.
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      name: Python `str`, name given to ops managed by this object. Default:
        E.g., `Blockwise([Exp(), Softplus()]).name ==
        'blockwise_of_exp_and_softplus'`.

    Raises:
      NotImplementedError: If a bijector with `event_ndims` > 1 or one that
        reshapes events is passed.
      ValueError: If `bijectors` list is empty.
      ValueError: If size of `block_sizes` does not equal to the length of
        bijectors or is not a vector.
    """
    super(Blockwise, self).__init__(
        forward_min_event_ndims=1,
        validate_args=validate_args,
        name=name or
        'blockwise_of_' + '_and_'.join([b.name for b in bijectors]))

    if not bijectors:
      raise ValueError('`bijectors` must not be empty.')

    for bijector in bijectors:
      if (bijector.forward_min_event_ndims > 1 or
          bijector.inverse_min_event_ndims != bijector.forward_min_event_ndims):
        # TODO(b/) In the future, it can be reasonable to support N-D bijectors
        # by concatenating along some specific axis, broadcasting low-D
        # bijectors appropriately.
        raise NotImplementedError('Only scalar and vector event-shape '
                                  'bijectors that do not alter the '
                                  'shape are supported at this time.')

    self._bijectors = bijectors

    with self._name_scope('init'):
      if block_sizes is None:
        block_sizes = tf.ones(len(bijectors), dtype=tf.int32)
      self._block_sizes = tf.convert_to_tensor(
          value=block_sizes, name='block_sizes', dtype_hint=tf.int32)

      self._block_sizes = _validate_block_sizes(self._block_sizes, bijectors,
                                                validate_args)

  @property
  def bijectors(self):
    return self._bijectors

  @property
  def block_sizes(self):
    return self._block_sizes

  def _forward(self, x):
    split_x = tf.split(x, self.block_sizes, axis=-1, num=len(self.bijectors))
    split_y = [b.forward(x_) for b, x_ in zip(self.bijectors, split_x)]
    y = tf.concat(split_y, axis=-1)
    tensorshape_util.set_shape(y, x.shape)
    return y

  def _inverse(self, y):
    split_y = tf.split(y, self.block_sizes, axis=-1, num=len(self.bijectors))
    split_x = [b.inverse(y_) for b, y_ in zip(self.bijectors, split_y)]
    x = tf.concat(split_x, axis=-1)
    tensorshape_util.set_shape(x, y.shape)
    return x

  def _forward_log_det_jacobian(self, x):
    split_x = tf.split(x, self.block_sizes, axis=-1, num=len(self.bijectors))
    fldjs = [
        b.forward_log_det_jacobian(x_, event_ndims=1)
        for b, x_ in zip(self.bijectors, split_x)
    ]
    return sum(fldjs)

  def _inverse_log_det_jacobian(self, y):
    split_y = tf.split(y, self.block_sizes, axis=-1, num=len(self.bijectors))
    ildjs = [
        b.inverse_log_det_jacobian(y_, event_ndims=1)
        for b, y_ in zip(self.bijectors, split_y)
    ]
    return sum(ildjs)

def _validate_block_sizes(block_sizes, bijectors, validate_args):
  """Helper to validate block sizes."""
  block_sizes_shape = block_sizes.shape
  if tensorshape_util.is_fully_defined(block_sizes_shape):
    if (tensorshape_util.rank(block_sizes_shape) != 1 or
        (tensorshape_util.num_elements(block_sizes_shape) != len(bijectors))):
      raise ValueError(
          '`block_sizes` must be `None`, or a vector of the same length as '
          '`bijectors`. Got a `Tensor` with shape {} and `bijectors` of '
          'length {}'.format(block_sizes_shape, len(bijectors)))
    return block_sizes
  elif validate_args:
    message = ('`block_sizes` must be `None`, or a vector of the same length '
               'as `bijectors`.')
    with tf.control_dependencies([
        assert_util.assert_equal(
            tf.size(input=block_sizes), len(bijectors), message=message),
        assert_util.assert_equal(tf.rank(block_sizes), 1)
    ]):
      return tf.identity(block_sizes)
  else:
    return block_sizes


# 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 `JointDistribution` base class."""

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

import abc
import six

import tensorflow.compat.v2 as tf

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

__all__ = [
    'JointDistribution',
]

@six.add_metaclass(abc.ABCMeta)
class JointDistribution(distribution_lib.Distribution):
  """Joint distribution over one or more component distributions.

  This distribution enables both sampling and joint probability computation from
  a single model specification.

  A joint distribution is a collection of possibly interdependent distributions.

  **Note**: unlike other non-`JointDistribution` distributions in
  `tfp.distributions`, `JointDistribution.sample` (and subclasses) return a
  structure of  `Tensor`s rather than a `Tensor`.  A structure can be a `list`,
  `tuple`, `dict`, `collections.namedtuple`, etc. Accordingly
  `joint.batch_shape` returns a structure of `TensorShape`s for each of the
  distributions' batch shapes and `joint.batch_shape_tensor()` returns a
  structure of `Tensor`s for each of the distributions' event shapes. (Same with
  `event_shape` analogues.)

  #### Subclass Requirements

  Subclasses implement:

  - `_flat_sample_distributions`: returns two `list`-likes: the first being a
    sequence of `Distribution`-like instances the second being a sequence of
    `Tensor` samples, each one drawn from its corresponding `Distribution`-like
    instance. The optional `value` argument is either `None` or a `list`-like
    with the same `len` as either of the results.

  - `_flatten`: takes a structured input and returns a sequence.

  - `_unflatten`: takes a sequence and returns a structure matching the
    semantics of the `JointDistribution` subclass.

  Subclasses initialize:

  - `_most_recently_built_distributions`: an iterable sequence of distributions
    which are known at `__init__` or `None`.
  """

  @property
  def model(self):
    return self._model

  @abc.abstractmethod
  def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):
    raise NotImplementedError()

  @abc.abstractmethod
  def _unflatten(self, xs):
    raise NotImplementedError()

  @abc.abstractmethod
  def _flatten(self, xs):
    raise NotImplementedError()

  @property
  def dtype(self):
    """The `DType` of `Tensor`s handled by this `Distribution`."""
    if self._most_recently_built_distributions is None:
      return super(JointDistribution, self).dtype
    return self._unflatten(
        None if d is None else d.dtype
        for d in self._most_recently_built_distributions)

  @property
  def reparameterization_type(self):
    """Describes how samples from the distribution are reparameterized.

    Currently this is one of the static instances
    `tfd.FULLY_REPARAMETERIZED` or `tfd.NOT_REPARAMETERIZED`.

    Returns:
      reparameterization_type: `ReparameterizationType` of each distribution in
        `model`.
    """
    if self._most_recently_built_distributions is None:
      return None
    return self._unflatten(
        None if d is None else d.reparameterization_type
        for d in self._most_recently_built_distributions)

  @property
  def batch_shape(self):
    """Shape of a single sample from a single event index as a `TensorShape`.

    May be partially defined or unknown.

    The batch dimensions are indexes into independent, non-identical
    parameterizations of this distribution.

    Returns:
      batch_shape: `tuple` of `TensorShape`s representing the `batch_shape` for
        each distribution in `model`.
    """
    # The following cannot leak graph Tensors in eager because `batch_shape` is
    # a `TensorShape`.
    if self._most_recently_built_distributions is None:
      return None
    return self._unflatten(
        tf.TensorShape(None) if d is None else d.batch_shape
        for d in self._most_recently_built_distributions)

  def batch_shape_tensor(self, name='batch_shape_tensor'):
    """Shape of a single sample from a single event index as a 1-D `Tensor`.

    The batch dimensions are indexes into independent, non-identical
    parameterizations of this distribution.

    Args:
      name: name to give to the op

    Returns:
      batch_shape: `Tensor` representing batch shape of each distribution in
        `model`.
    """
    with self._name_scope(name):
      return self._unflatten(self._map_attr_over_dists('batch_shape_tensor')[1])

  @property
  def event_shape(self):
    """Shape of a single sample from a single batch as a `TensorShape`.

    May be partially defined or unknown.

    Returns:
      event_shape: `tuple` of `TensorShape`s representing the `event_shape` for
        each distribution in `model`.
    """
    # The following cannot leak graph Tensors in eager because `batch_shape` is
    # a `TensorShape`.
    if self._most_recently_built_distributions is None:
      return None
    return self._unflatten(
        tf.TensorShape(None) if d is None else d.event_shape
        for d in self._most_recently_built_distributions)

  def event_shape_tensor(self, name='event_shape_tensor'):
    """Shape of a single sample from a single batch as a 1-D int32 `Tensor`.

    Args:
      name: name to give to the op

    Returns:
      event_shape: `tuple` of `Tensor`s representing the `event_shape` for each
        distribution in `model`.
    """
    with self._name_scope(name):
      return self._unflatten(self._map_attr_over_dists('event_shape_tensor')[1])

  def sample_distributions(self, sample_shape=(), seed=None, value=None,
                           name='sample_distributions'):
    """Generate samples and the (random) distributions.

    Note that a call to `sample()` without arguments will generate a single
    sample.

    Args:
      sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples.
      seed: Python integer seed for generating random numbers.
      value: `list` of `Tensor`s in `distribution_fn` order to use to
        parameterize other ("downstream") distribution makers.
        Default value: `None` (i.e., draw a sample from each distribution).
      name: name prepended to ops created by this function.
        Default value: `"sample_distributions"`.

    Returns:
      distributions: a `tuple` of `Distribution` instances for each of
        `distribution_fn`.
      samples: a `tuple` of `Tensor`s with prepended dimensions `sample_shape`
        for each of `distribution_fn`.
    """
    with self._name_scope(name):
      ds, xs = self._call_flat_sample_distributions(sample_shape, seed, value)
      return self._unflatten(ds), self._unflatten(xs)

  def log_prob_parts(self, value, name='log_prob_parts'):
    """Log probability density/mass function.

    Args:
      value: `list` of `Tensor`s in `distribution_fn` order for which we compute
        the `log_prob_parts` and to parameterize other ("downstream")
        distributions.
      name: name prepended to ops created by this function.
        Default value: `"log_prob_parts"`.

    Returns:
      log_prob_parts: a `tuple` of `Tensor`s representing the `log_prob` for
        each `distribution_fn` evaluated at each corresponding `value`.
    """
    with self._name_scope(name):
      _, xs = self._map_measure_over_dists('log_prob', value)
      return self._unflatten(xs)

  def prob_parts(self, value, name='prob_parts'):
    """Log probability density/mass function.

    Args:
      value: `list` of `Tensor`s in `distribution_fn` order for which we compute
        the `prob_parts` and to parameterize other ("downstream") distributions.
      name: name prepended to ops created by this function.
        Default value: `"prob_parts"`.

    Returns:
      prob_parts: a `tuple` of `Tensor`s representing the `prob` for
        each `distribution_fn` evaluated at each corresponding `value`.
    """
    with self._name_scope(name):
      _, xs = self._map_measure_over_dists('prob', value)
      return self._unflatten(xs)

  def is_scalar_event(self, name='is_scalar_event'):
    """Indicates that `event_shape == []`.

    Args:
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      is_scalar_event: `bool` scalar `Tensor` for each distribution in `model`.
    """
    with self._name_scope(name):
      return self._unflatten(self._map_attr_over_dists('is_scalar_event')[1])

  def is_scalar_batch(self, name='is_scalar_batch'):
    """Indicates that `batch_shape == []`.

    Args:
      name: Python `str` prepended to names of ops created by this function.

    Returns:
      is_scalar_batch: `bool` scalar `Tensor` for each distribution in `model`.
    """
    with self._name_scope(name):
      return self._unflatten(self._map_attr_over_dists('is_scalar_batch')[1])

  def _log_prob(self, value):
    _, xs = self._map_measure_over_dists('log_prob', value)
    return sum(xs)

  @distribution_util.AppendDocstring(kwargs_dict={
      'value': ('`Tensor`s structured like `type(model)` used to parameterize '
                'other dependent ("downstream") distribution-making functions. '
                'Using `None` for any element will trigger a sample from the '
                'corresponding distribution. Default value: `None` '
                '(i.e., draw a sample from each distribution).')})
  def _sample_n(self, sample_shape, seed, value=None):
    _, xs = self._call_flat_sample_distributions(sample_shape, seed, value)
    return self._unflatten(xs)

  def _map_measure_over_dists(self, attr, value):
    if any(x is None for x in tf.nest.flatten(value)):
      raise ValueError('No `value` part can be `None`; saw: {}.'.format(value))
    ds, xs = self._call_flat_sample_distributions(value=value)
    return ds, maybe_check_wont_broadcast(
        (getattr(d, attr)(x) for d, x in zip(ds, xs)),
        self.validate_args)

  def _map_attr_over_dists(self, attr):
    if (not self._most_recently_built_distributions or
        any(d is None for d in self._most_recently_built_distributions)):
      # Same seed to help CSE.
      ds, _ = self._call_flat_sample_distributions(seed=42)
    else:
      ds = self._most_recently_built_distributions
    return ds, (None if d is None else getattr(d, attr)() for d in ds)

  def _call_flat_sample_distributions(
      self, sample_shape=(), seed=None, value=None):
    if value is not None:
      value = self._flatten(value)
    ds, xs = self._flat_sample_distributions(sample_shape, seed, value)
    if not value:
      self._most_recently_built_distributions = tuple(ds)
    return ds, xs

  # We need to bypass base Distribution reshaping/validation logic so we
  # tactically implement a few of the `_call_*` redirectors.  We don't want to
  # override the public level because then tfp.layers can't take generic
  # `Distribution.foo` as argument for the `convert_to_tensor_fn` parameter.
  def _call_log_prob(self, value, name):
    with self._name_scope(name):
      return self._log_prob(value)

  def _call_sample_n(self, sample_shape, seed, name, value=None):
    with self._name_scope(name):
      return self._sample_n(sample_shape, seed, value)

def maybe_check_wont_broadcast(flat_xs, validate_args):
  """Verifies that `parts` don't broadcast."""
  flat_xs = tuple(flat_xs)  # So we can receive generators.
  if not validate_args:
    # Note: we don't try static validation because it is theoretically
    # possible that a user wants to take advantage of broadcasting.
    # Only when `validate_args` is `True` do we enforce the validation.
    return flat_xs
  msg = 'Broadcasting probably indicates an error in model specification.'
  s = tuple(x.shape for x in flat_xs)
  if all(tensorshape_util.is_fully_defined(s_) for s_ in s):
    if not all(a == b for a, b in zip(s[1:], s[:-1])):
      raise ValueError(msg)
    return flat_xs
  assertions = [assert_util.assert_equal(a, b, message=msg)
                for a, b in zip(s[1:], s[:-1])]
  with tf.control_dependencies(assertions):
    return tuple(tf.identity(x) for x in flat_xs)


# 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.
# ============================================================================
"""Trainable distributions.

"Trainable distributions" are instances of `tfp.distributions` which are
parameterized by a transformation of a single input `Tensor`. The
transformations are presumed to use TensorFlow variables and typically need to
be fit, e.g., using `tf.train` optimizers or `tfp.optimizers`.
"""

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

import tensorflow as tf

from tensorflow_probability.python import distributions as tfd

__all__ = [
    'bernoulli',
    'multivariate_normal_tril',
    'normal',
    'poisson',
    'softplus_and_shift',
    'tril_with_diag_softplus_and_shift',
]

def softplus_and_shift(x, shift=1e-5, name=None):
  """Converts (batch of) scalars to (batch of) positive valued scalars.

  Args:
    x: (Batch of) `float`-like `Tensor` representing scalars which will be
      transformed into positive elements.
    shift: `Tensor` added to `softplus` transformation of elements.
      Default value: `1e-5`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "positive_tril_with_shift").

  Returns:
    scale: (Batch of) scalars`with `x.dtype` and `x.shape`.
  """
  with tf.compat.v1.name_scope(name, 'softplus_and_shift', [x, shift]):
    x = tf.convert_to_tensor(value=x, name='x')
    y = tf.nn.softplus(x)
    if shift is not None:
      y += shift
    return y

def tril_with_diag_softplus_and_shift(x, diag_shift=1e-5, name=None):
  """Converts (batch of) vectors to (batch of) lower-triangular scale matrices.

  Args:
    x: (Batch of) `float`-like `Tensor` representing vectors which will be
      transformed into lower-triangular scale matrices with positive diagonal
      elements. Rightmost shape `n` must be such that
      `n = dims * (dims + 1) / 2` for some positive, integer `dims`.
    diag_shift: `Tensor` added to `softplus` transformation of diagonal
      elements.
      Default value: `1e-5`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "tril_with_diag_softplus_and_shift").

  Returns:
    scale_tril: (Batch of) lower-triangular `Tensor` with `x.dtype` and
      rightmost shape `[dims, dims]` where `n = dims * (dims + 1) / 2` where
      `n = x.shape[-1]`.
  """
  with tf.compat.v1.name_scope(name, 'tril_with_diag_softplus_and_shift',
                               [x, diag_shift]):
    x = tf.convert_to_tensor(value=x, name='x')
    x = tfd.fill_triangular(x)
    diag = softplus_and_shift(tf.linalg.diag_part(x), diag_shift)
    x = tf.linalg.set_diag(x, diag)
    return x

def multivariate_normal_tril(x,
                             dims,
                             layer_fn=tf.compat.v1.layers.dense,
                             loc_fn=lambda x: x,
                             scale_fn=tril_with_diag_softplus_and_shift,
                             name=None):
  """Constructs a trainable `tfd.MultivariateNormalTriL` distribution.

  This function creates a MultivariateNormal (MVN) with lower-triangular scale
  matrix. By default the MVN is parameterized via affine transformation of input
  tensor `x`. Using default args, this function is mathematically equivalent to:

  ```none
  Y = MVN(loc=matmul(W, x) + b,
          scale_tril=f(reshape_tril(matmul(M, x) + c)))

  where,
    W in R^[d, n]
    M in R^[d*(d+1)/2, n]
    b in R^d
    c in R^d
    f(S) = set_diag(S, softplus(matrix_diag_part(S)) + 1e-5)
  ```

  Observe that `f` makes the diagonal of the triangular-lower scale matrix be
  positive and no smaller than `1e-5`.

  #### Examples

  ```python
  # This example fits a multilinear regression loss.
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  y_size = 2  # size of single y
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size, y_size).astype(dtype)
    b = np.random.randn(1, y_size).astype(dtype)
    true_mean = np.tensordot(x, w, axes=[[-1], [0]]) + b
    noise = np.random.randn(n, y_size).astype(dtype)
    y = true_mean + noise
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting MVNTriL maximum likelihood estimator.
  mvn = tfp.trainable_distributions.multivariate_normal_tril(x, dims=y_size)
  loss = -tf.reduce_mean(mvn.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-3).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, mvn.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:38.2020797729  mse:4.17175960541
  #     iteration:200  loss:2.90179634094  mse:0.990987896919
  #     iteration:400  loss:2.82727336884  mse:0.990926623344
  #     iteration:600  loss:2.82726788521  mse:0.990926682949
  #     iteration:800  loss:2.82726788521  mse:0.990926682949
  #     iteration:999  loss:2.82726788521  mse:0.990926682949
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    dims: Scalar, `int`, `Tensor` indicated the MVN event size, i.e., the
      created MVN will be distribution over length-`dims` vectors.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [d]], axis=0)`.
      Default value: `tf.layers.dense`.
    loc_fn: Python `callable` which transforms the `loc` parameter. Takes a
      (batch of) length-`dims` vectors and returns a `Tensor` of same shape and
      `dtype`.
      Default value: `lambda x: x`.
    scale_fn: Python `callable` which transforms the `scale` parameters. Takes a
      (batch of) length-`dims * (dims + 1) / 2` vectors and returns a
      lower-triangular `Tensor` of same batch shape with rightmost dimensions
      having shape `[dims, dims]`.
      Default value: `tril_with_diag_softplus_and_shift`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "multivariate_normal_tril").

  Returns:
    mvntril: An instance of `tfd.MultivariateNormalTriL`.
  """
  with tf.compat.v1.name_scope(name, 'multivariate_normal_tril', [x, dims]):
    x = tf.convert_to_tensor(value=x, name='x')
    x = layer_fn(x, dims + dims * (dims + 1) // 2)
    return tfd.MultivariateNormalTriL(
        loc=loc_fn(x[..., :dims]),
        scale_tril=scale_fn(x[..., dims:]))

def bernoulli(x, layer_fn=tf.compat.v1.layers.dense, name=None):
  """Constructs a trainable `tfd.Bernoulli` distribution.

  This function creates a Bernoulli distribution parameterized by logits.
  Using default args, this function is mathematically equivalent to:

  ```none
  Y = Bernoulli(logits=matmul(W, x) + b)

  where,
    W in R^[d, n]
    b in R^d
  ```

  #### Examples

  This function can be used as a [logistic regression](
  https://en.wikipedia.org/wiki/Logistic_regression) loss.

  ```python
  # This example fits a logistic regression loss.
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size).astype(dtype)
    b = np.random.randn(1).astype(dtype)
    true_logits = np.tensordot(x, w, axes=[[-1], [-1]]) + b
    noise = np.random.logistic(size=n).astype(dtype)
    y = dtype(true_logits + noise > 0.)
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting Bernoulli maximum likelihood estimator.
  bernoulli = tfp.trainable_distributions.bernoulli(x)
  loss = -tf.reduce_mean(bernoulli.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, bernoulli.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:0.635675370693  mse:0.222526371479
  #     iteration:200  loss:0.440077394247  mse:0.143687799573
  #     iteration:400  loss:0.440077394247  mse:0.143687844276
  #     iteration:600  loss:0.440077394247  mse:0.143687844276
  #     iteration:800  loss:0.440077424049  mse:0.143687844276
  #     iteration:999  loss:0.440077424049  mse:0.143687844276
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
      Default value: `tf.layers.dense`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "bernoulli").

  Returns:
    bernoulli: An instance of `tfd.Bernoulli`.
  """
  with tf.compat.v1.name_scope(name, 'bernoulli', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    logits = tf.squeeze(layer_fn(x, 1), axis=-1)
    return tfd.Bernoulli(logits=logits)

def normal(x,
           layer_fn=tf.compat.v1.layers.dense,
           loc_fn=lambda x: x,
           scale_fn=1.,
           name=None):
  """Constructs a trainable `tfd.Normal` distribution.

  This function creates a Normal distribution parameterized by loc and scale.
  Using default args, this function is mathematically equivalent to:

  ```none
  Y = Normal(loc=matmul(W, x) + b, scale=1)

  where,
    W in R^[d, n]
    b in R^d
  ```

  #### Examples

  This function can be used as a [linear regression](
  https://en.wikipedia.org/wiki/Linear_regression) loss.

  ```python
  # This example fits a linear regression loss.
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size).astype(dtype)
    b = np.random.randn(1).astype(dtype)
    true_mean = np.tensordot(x, w, axes=[[-1], [-1]]) + b
    noise = np.random.randn(n).astype(dtype)
    y = true_mean + noise
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting Normal maximum likelihood estimator.
  normal = tfp.trainable_distributions.normal(x)
  loss = -tf.reduce_mean(normal.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, normal.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:6.34114170074  mse:10.8444051743
  #     iteration:200  loss:1.40146839619  mse:0.965059816837
  #     iteration:400  loss:1.40052902699  mse:0.963181257248
  #     iteration:600  loss:1.40052902699  mse:0.963181257248
  #     iteration:800  loss:1.40052902699  mse:0.963181257248
  #     iteration:999  loss:1.40052902699  mse:0.963181257248
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
      Default value: `tf.layers.dense`.
    loc_fn: Python `callable` which transforms the `loc` parameter. Takes a
      (batch of) length-`dims` vectors and returns a `Tensor` of same shape and
      `dtype`.
      Default value: `lambda x: x`.
    scale_fn: Python `callable` or `Tensor`. If a `callable` transforms the
      `scale` parameters; if `Tensor` is the `tfd.Normal` `scale` argument.
      Takes a (batch of) length-`dims` vectors and returns a `Tensor` of same
      size. (Taking a `callable` or `Tensor` is how `tf.Variable` intializers
      behave.)
      Default value: `1`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "normal").

  Returns:
    normal: An instance of `tfd.Normal`.
  """
  with tf.compat.v1.name_scope(name, 'normal', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    if callable(scale_fn):
      y = layer_fn(x, 2)
      loc = loc_fn(y[..., 0])
      scale = scale_fn(y[..., 1])
    else:
      y = tf.squeeze(layer_fn(x, 1), axis=-1)
      loc = loc_fn(y)
      scale = tf.cast(scale_fn, loc.dtype.base_dtype)
    return tfd.Normal(loc=loc, scale=scale)

def poisson(x,
            layer_fn=tf.compat.v1.layers.dense,
            log_rate_fn=lambda x: x,
            name=None):
  """Constructs a trainable `tfd.Poisson` distribution.

  This function creates a Poisson distribution parameterized by log rate.
  Using default args, this function is mathematically equivalent to:

  ```none
  Y = Poisson(log_rate=matmul(W, x) + b)

  where,
    W in R^[d, n]
    b in R^d
  ```

  #### Examples

  This can be used as a [Poisson regression](
  https://en.wikipedia.org/wiki/Poisson_regression) loss.

  ```python
  # This example fits a poisson regression loss.
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size).astype(dtype)
    b = np.random.randn(1).astype(dtype)
    true_log_rate = np.tensordot(x, w, axes=[[-1], [-1]]) + b
    y = np.random.poisson(lam=np.exp(true_log_rate)).astype(dtype)
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting Poisson maximum likelihood estimator.
  poisson = tfp.trainable_distributions.poisson(x)
  loss = -tf.reduce_mean(poisson.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, poisson.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:37.0814208984  mse:6359.41259766
  #     iteration:200  loss:1.42010736465  mse:40.7654914856
  #     iteration:400  loss:1.39027583599  mse:8.77660560608
  #     iteration:600  loss:1.3902695179   mse:8.78443241119
  #     iteration:800  loss:1.39026939869  mse:8.78443622589
  #     iteration:999  loss:1.39026939869  mse:8.78444766998
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
      Default value: `tf.layers.dense`.
    log_rate_fn: Python `callable` which transforms the `log_rate` parameter.
      Takes a (batch of) length-`dims` vectors and returns a `Tensor` of same
      shape and `dtype`.
      Default value: `lambda x: x`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "poisson").

  Returns:
    poisson: An instance of `tfd.Poisson`.
  """
  with tf.compat.v1.name_scope(name, 'poisson', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    log_rate = log_rate_fn(tf.squeeze(layer_fn(x, 1), axis=-1))
    return tfd.Poisson(log_rate=log_rate)


# 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.
# ============================================================================
"""Methods for finding roots of functions of one variable."""

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

import collections

import tensorflow as tf

__all__ = [
    'secant_root',
]

RootSearchResults = collections.namedtuple(
    'RootSearchResults',
    [
        # A tensor containing the last position explored. If the search was
        # successful, this position is a root of the objective function.
        'estimated_root',
        # A tensor containing the value of the objective function at the last
        # position explored. If the search was successful, then this is close
        # to 0.
        'objective_at_estimated_root',
        # The number of iterations performed.
        'num_iterations',
    ])

def secant_root(objective_fn,
                initial_position,
                next_position=None,
                value_at_position=None,
                position_tolerance=1e-8,
                value_tolerance=1e-8,
                max_iterations=50,
                stopping_policy_fn=tf.reduce_all,
                validate_args=False,
                name=None):
  r"""Finds root(s) of a function of single variable using the secant method.

  The [secant method](https://en.wikipedia.org/wiki/Secant_method) is a
  root-finding algorithm that uses a succession of roots of secant lines to
  better approximate a root of a function. The secant method can be thought of
  as a finite-difference approximation of Newton's method.

  Args:
    objective_fn: Python callable for which roots are searched. It must be a
      callable of a single variable. `objective_fn` must return a `Tensor` of
      the same shape and dtype as `initial_position`.
    initial_position: `Tensor` or Python float representing the starting
      position. The function will search for roots in the neighborhood of each
      point. The shape of `initial_position` should match that of the input to
      `objective_fn`.
    next_position: Optional `Tensor` representing the next position in the
      search. If specified, this argument must broadcast with the shape of
      `initial_position` and have the same dtype. It will be used to compute the
      first step to take when searching for roots. If not specified, a default
      value will be used instead.
      Default value: `initial_position * (1 + 1e-4) + sign(initial_position) *
        1e-4`.
    value_at_position: Optional `Tensor` or Pyhon float representing the value
      of `objective_fn` at `initial_position`. If specified, this argument must
      have the same shape and dtype as `initial_position`. If not specified, the
      value will be evaluated during the search.
      Default value: None.
    position_tolerance: Optional `Tensor` representing the tolerance for the
      estimated roots. If specified, this argument must broadcast with the shape
      of `initial_position` and have the same dtype.
      Default value: `1e-8`.
    value_tolerance: Optional `Tensor` representing the tolerance used to check
      for roots. If the absolute value of `objective_fn` is smaller than
      `value_tolerance` at a given position, then that position is considered a
      root for the function. If specified, this argument must broadcast with the
      shape of `initial_position` and have the same dtype.
      Default value: `1e-8`.
    max_iterations: Optional `Tensor` or Python integer specifying the maximum
      number of steps to perform for each initial position. Must broadcast with
      the shape of `initial_position`.
      Default value: `50`.
    stopping_policy_fn: Python `callable` controlling the algorithm termination.
      It must be a callable accepting a `Tensor` of booleans with the shape of
      `initial_position` (each denoting whether the search is finished for each
      starting point), and returning a scalar boolean `Tensor` (indicating
      whether the overall search should stop). Typical values are
      `tf.reduce_all` (which returns only when the search is finished for all
      points), and `tf.reduce_any` (which returns as soon as the search is
      finished for any point).
      Default value: `tf.reduce_all` (returns only when the search is finished
        for all points).
    validate_args: Python `bool` indicating whether to validate arguments such
      as `position_tolerance`, `value_tolerance`, and `max_iterations`.
      Default value: `False`.
    name: Python `str` name prefixed to ops created by this function.

  Returns:
    root_search_results: A Python `namedtuple` containing the following items:
      estimated_root: `Tensor` containing the last position explored. If the
        search was successful within the specified tolerance, this position is
        a root of the objective function.
      objective_at_estimated_root: `Tensor` containing the value of the
        objective function at `position`. If the search was successful within
        the specified tolerance, then this is close to 0.
      num_iterations: The number of iterations performed.

  Raises:
    ValueError: if a non-callable `stopping_policy_fn` is passed.

  #### Examples

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tf.enable_eager_execution()

  # Example 1: Roots of a single function from two different starting points.

  f = lambda x: (63 * x**5 - 70 * x**3 + 15 * x) / 8.
  x = tf.constant([-1, 10], dtype=tf.float64)

  tfp.math.secant_root(objective_fn=f, initial_position=x))
  # ==> RootSearchResults(
      estimated_root=array([-0.90617985, 0.90617985]),
      objective_at_estimated_root=array([-4.81727769e-10, 7.44957651e-10]),
      num_iterations=array([ 7, 24], dtype=int32))

  tfp.math.secant_root(objective_fn=f,
                       initial_position=x,
                       stopping_policy_fn=tf.reduce_any)
  # ==> RootSearchResults(
      estimated_root=array([-0.90617985, 3.27379206]),
      objective_at_estimated_root=array([-4.81727769e-10, 2.66058312e+03]),
      num_iterations=array([7, 8], dtype=int32))

  # Example 2: Roots of a multiplex function from a single starting point.

  def f(x):
    return tf.constant([0., 63. / 8], dtype=tf.float64) * x**5 \
        + tf.constant([5. / 2, -70. / 8], dtype=tf.float64) * x**3 \
        + tf.constant([-3. / 2, 15. / 8], dtype=tf.float64) * x

  x = tf.constant([-1, -1], dtype=tf.float64)

  tfp.math.secant_root(objective_fn=f, initial_position=x)
  # ==> RootSearchResults(
      estimated_root=array([-0.77459667, -0.90617985]),
      objective_at_estimated_root=array([-7.81339438e-11, -4.81727769e-10]),
      num_iterations=array([7, 7], dtype=int32))

  # Example 3: Roots of a multiplex function from two starting points.

  def f(x):
    return tf.constant([0., 63. / 8], dtype=tf.float64) * x**5 \
        + tf.constant([5. / 2, -70. / 8], dtype=tf.float64) * x**3 \
        + tf.constant([-3. / 2, 15. / 8], dtype=tf.float64) * x

  x = tf.constant([[-1, -1], [10, 10]], dtype=tf.float64)

  tfp.math.secant_root(objective_fn=f, initial_position=x)
  # ==> RootSearchResults(
      estimated_root=array([
          [-0.77459667, -0.90617985],
          [ 0.77459667, 0.90617985]]),
      objective_at_estimated_root=array([
          [-7.81339438e-11, -4.81727769e-10],
          [6.66025013e-11, 7.44957651e-10]]),
      num_iterations=array([
          [7, 7],
          [16, 24]], dtype=int32))
  ```
  """
  if not callable(stopping_policy_fn):
    raise ValueError('stopping_policy_fn must be callable')

  position = tf.convert_to_tensor(
      value=initial_position,
      name='position',
  )
  value_at_position = tf.convert_to_tensor(
      value=value_at_position or objective_fn(position),
      name='value_at_position',
      dtype=position.dtype.base_dtype)

  zero = tf.zeros_like(position)
  position_tolerance = tf.convert_to_tensor(
      value=position_tolerance, name='position_tolerance', dtype=position.dtype)
  value_tolerance = tf.convert_to_tensor(
      value=value_tolerance, name='value_tolerance', dtype=position.dtype)

  num_iterations = tf.zeros_like(position, dtype=tf.int32)
  max_iterations = tf.convert_to_tensor(value=max_iterations, dtype=tf.int32)
  max_iterations = tf.broadcast_to(
      max_iterations, name='max_iterations', shape=position.shape)

  # Compute the step from `next_position` if present. This covers the case where
  # a user has two starting points, which bound the root or has a specific step
  # size in mind.
  if next_position is None:
    epsilon = tf.constant(1e-4, dtype=position.dtype, shape=position.shape)
    step = position * epsilon + tf.sign(position) * epsilon
  else:
    step = next_position - initial_position

  finished = tf.constant(False, shape=position.shape)

  # Negate `stopping_condition` to determine if the search should continue.
  # This means, in particular, that tf.reduce_*all* will return only when the
  # search is finished for *all* starting points.
  def _should_continue(position, value_at_position, num_iterations, step,
                       finished):
    """Indicates whether the overall search should continue.

    Args:
      position: `Tensor` containing the current root estimates.
      value_at_position: `Tensor` containing the value of `objective_fn` at
        `position`.
      num_iterations: `Tensor` containing the current iteration index for each
        point.
      step: `Tensor` containing the size of the step to take for each point.
      finished: `Tensor` indicating for which points the search is finished.

    Returns:
      A boolean value indicating whether the overall search should continue.
    """
    del position, value_at_position, num_iterations, step  # Unused
    return ~tf.convert_to_tensor(
        value=stopping_policy_fn(finished), name='should_stop', dtype=tf.bool)

  # For each point in `position`, the search is stopped if either:
  # (1) A root has been found
  # (2) f(position) == f(position + step)
  # (3) The maximum number of iterations has been reached
  # In case (2), the search may be stopped both before the desired tolerance is
  # achieved (or even a root is found), and the maximum number of iterations is
  # reached.
  def _body(position, value_at_position, num_iterations, step, finished):
    """Performs one iteration of the secant root-finding algorithm.

    Args:
      position: `Tensor` containing the current root estimates.
      value_at_position: `Tensor` containing the value of `objective_fn` at
        `position`.
      num_iterations: `Tensor` containing the current iteration index for each
        point.
      step: `Tensor` containing the size of the step to take for each point.
      finished: `Tensor` indicating for which points the search is finished.

    Returns:
      The `Tensor`s to use for the next iteration of the algorithm.
    """

    # True if the search was already finished, or (1) or (3) just became true.
    was_finished = finished | (num_iterations >= max_iterations) | (
        tf.abs(step) < position_tolerance) | (
            tf.abs(value_at_position) < value_tolerance)

    # Compute the next position and the value at that point.
    next_position = tf.where(was_finished, position, position + step)
    value_at_next_position = tf.where(was_finished, value_at_position,
                                      objective_fn(next_position))

    # True if the search was already finished, or (2) just became true.
    is_finished = tf.equal(value_at_position, value_at_next_position)

    # Use the mid-point between the last two positions if (2) just became true.
    next_position = tf.where(is_finished & ~was_finished,
                             (position + next_position) * 0.5, next_position)

    # Once finished, stop updating the iteration index and set the step to zero.
    num_iterations = tf.where(is_finished, num_iterations, num_iterations + 1)
    next_step = tf.where(
        is_finished, zero, step * value_at_next_position /
        (value_at_position - value_at_next_position))

    return (next_position, value_at_next_position, num_iterations, next_step,
            is_finished)

  with tf.compat.v1.name_scope(
      name, 'secant_root',
      [position, next_position, value_at_position, max_iterations]):

    assertions = []
    if validate_args:
      assertions += [
          tf.Assert(
              tf.reduce_all(input_tensor=position_tolerance > zero),
              [position_tolerance]),
          tf.Assert(
              tf.reduce_all(input_tensor=value_tolerance > zero),
              [value_tolerance]),
          tf.Assert(
              tf.reduce_all(input_tensor=max_iterations >= num_iterations),
              [max_iterations]),
      ]

    with tf.control_dependencies(assertions):
      root, value_at_root, num_iterations, _, _ = tf.while_loop(
          cond=_should_continue,
          body=_body,
          loop_vars=[
              position, value_at_position, num_iterations, step, finished
          ])

  return RootSearchResults(
      estimated_root=root,
      objective_at_estimated_root=value_at_root,
      num_iterations=num_iterations)


# 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.
# ============================================================================
"""Metropolis-adjusted Langevin algorithm, a gradient-based MCMC algorithm."""

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.internal import distribution_util
from tensorflow_probability.python.math import diag_jacobian
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__ = [
    'MetropolisAdjustedLangevinAlgorithm',
    'UncalibratedLangevin',
]

UncalibratedLangevinKernelResults = collections.namedtuple(
    'UncalibratedLangevinKernelResults',
    [
        'grads_target_log_prob',
        # Results are for "next_state".
        'log_acceptance_correction',
        'target_log_prob',
        'volatility',
        'grads_volatility',
        'diffusion_drift',
    ])

class MetropolisAdjustedLangevinAlgorithm(kernel_base.TransitionKernel):
  """Runs one step of Metropolis-adjusted Langevin algorithm.

  Metropolis-adjusted Langevin algorithm (MALA) is a Markov chain Monte Carlo
  (MCMC) algorithm that takes a step of a discretised Langevin diffusion as a
  proposal. This class implements one step of MALA using Euler-Maruyama method
  for a given `current_state` and diagonal preconditioning `volatility` matrix.
  Mathematical details and derivations can be found in
  [Roberts and Rosenthal (1998)][1] and [Xifara et al. (2013)][2].

  See `UncalibratedLangevin` class description below for details on the proposal
  generating step of the algorithm.

  The `one_step` function 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 reduce 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:

  ##### Simple chain with warm-up.

  In this example we sample from a standard univariate normal
  distribution using MALA with `step_size` equal to 0.75.

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

  tfd = tfp.distributions

  dtype = np.float32

  with tf.Session(graph=tf.Graph()) as sess:
    # Target distribution is Standard Univariate Normal
    target = tfd.Normal(loc=dtype(0), scale=dtype(1))

    # Define MALA sampler with `step_size` equal to 0.75
    samples, _ = tfp.mcmc.sample_chain(
        num_results=1000,
        current_state=dtype(1),
        kernel=tfp.mcmc.MetropolisAdjustedLangevinAlgorithm(
            target_log_prob_fn=target.log_prob,
            step_size=0.75,
            seed=42),
        num_burnin_steps=500,
        parallel_iterations=1)  # For determinism.

    sample_mean = tf.reduce_mean(samples, axis=0)
    sample_std = tf.sqrt(
        tf.reduce_mean(tf.squared_difference(samples, sample_mean),
                       axis=0))

    sess.graph.finalize()  # No more graph building.

    [sample_mean_, sample_std_] = sess.run([sample_mean, sample_std])

  print('sample mean', sample_mean_)
  print('sample standard deviation', sample_std_)
  ```

  ##### Same example but in eager mode.

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  import numpy as np
  import matplotlib.pyplot as plt

  # Support for eager execution
  tf.enable_eager_execution()

  tfd = tfp.distributions
  dtype = np.float32

  # Target distribution is Standard Univariate Normal
  target = tfd.Normal(loc=dtype(0), scale=dtype(1))

  def target_log_prob(x):
    return target.log_prob(x)

  # Define MALA sampler with `step_size` equal to 0.75
  samples, _ = tfp.mcmc.sample_chain(
      num_results=1000,
      current_state=dtype(1),
      kernel=tfp.mcmc.MetropolisAdjustedLangevinAlgorithm(
          target_log_prob_fn=target_log_prob,
          step_size=0.75,
          seed=42),
      num_burnin_steps=500,
      parallel_iterations=1)  # For determinism.

  sample_mean = tf.reduce_mean(samples, axis=0)
  sample_std = tf.sqrt(
      tf.reduce_mean(tf.squared_difference(samples, sample_mean),
                     axis=0))

  print('sample mean', sample_mean)
  print('sample standard deviation', sample_std)

  plt.title('Traceplot')
  plt.plot(samples.numpy(), 'b')
  plt.xlabel('Iteration')
  plt.ylabel('Position')
  plt.show()
  ```

  ##### Sample from a 3-D Multivariate Normal distribution.

  In this example we also consider a non-constant volatility function.

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

  tfd = tfp.distributions

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

  with tf.Session(graph=tf.Graph()) as sess:
    # Target distribution is defined through the Cholesky decomposition
    chol = tf.linalg.cholesky(true_cov)
    target = tfd.MultivariateNormalTriL(loc=true_mean, scale_tril=chol)

    # Assume that the state is passed as a list of 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.concat([x, y], axis=-1) - true_mean
      return target.log_prob(z)

    # Here we define the volatility function to be non-constant
    def volatility_fn(x, y):
      # Stack the input tensors together
      return [1. / (0.5 + 0.1 * tf.sqrt(x * x)),
              1. / (0.5 + 0.1 *tf.sqrt(y * y))]

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

    # Run MALA with normal proposal for `num_results` iterations for
    # `num_chains` independent chains:
    states, _ = tfp.mcmc.sample_chain(
        num_results=num_results,
        current_state=init_state,
        kernel=tfp.mcmc.MetropolisAdjustedLangevinAlgorithm(
            target_log_prob_fn=target_log_prob,
            step_size=.1,
            volatility_fn=volatility_fn,
            seed=42),
        num_burnin_steps=200,
        num_steps_between_results=1,
        parallel_iterations=1)

    states = tf.concat(states, axis=-1)
    sample_mean = tf.reduce_mean(states, axis=[0, 1])
    x = tf.expand_dims(states - sample_mean, -1)
    sample_cov = tf.reduce_mean(
        tf.matmul(x, tf.transpose(x, [0, 1, 3, 2])), [0, 1])

    [sample_mean_, sample_cov_] = sess.run([
        sample_mean, sample_cov])

  print('sample mean', sample_mean_)
  print('sample covariance matrix', sample_cov_)
  ```

  #### References

  [1]: Gareth Roberts and Jeffrey Rosenthal. Optimal Scaling of Discrete
       Approximations to Langevin Diffusions. _Journal of the Royal Statistical
       Society: Series B (Statistical Methodology)_, 60: 255-268, 1998.
       https://doi.org/10.1111/1467-9868.00123

  [2]: T. Xifara et al. Langevin diffusions and the Metropolis-adjusted
       Langevin algorithm. _arXiv preprint arXiv:1309.2983_, 2013.
       https://arxiv.org/abs/1309.2983
  """

  def __init__(self,
               target_log_prob_fn,
               step_size,
               volatility_fn=None,
               seed=None,
               parallel_iterations=10,
               name=None):
    """Initializes MALA 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.
      step_size: `Tensor` or Python `list` of `Tensor`s representing the step
        size for the leapfrog integrator. Must broadcast with the shape of
        `current_state`. Larger step sizes lead to faster progress, but
        too-large step sizes make rejection exponentially more likely. When
        possible, it's often helpful to match per-variable step sizes to the
        standard deviations of the target distribution in each variable.
      volatility_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns
        volatility value at `current_state`. Should return a `Tensor` or Python
        `list` of `Tensor`s that must broadcast with the shape of
        `current_state` Defaults to the identity function.
      seed: Python integer to seed the random number generator.
      parallel_iterations: the number of coordinates for which the gradients of
        the volatility matrix `volatility_fn` can be computed in parallel.
        Default value: `None` (i.e., no seed).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'mala_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 `step_size` or a list with same length as
        `current_state`.
      TypeError: if `volatility_fn` is not callable.
    """
    impl = metropolis_hastings.MetropolisHastings(
        inner_kernel=UncalibratedLangevin(
            target_log_prob_fn=target_log_prob_fn,
            step_size=step_size,
            volatility_fn=volatility_fn,
            seed=seed,
            parallel_iterations=parallel_iterations,
            name=name),
        seed=seed)

    self._impl = impl
    parameters = impl.inner_kernel.parameters.copy()
    # Remove `compute_acceptance` parameter as this is not a MALA kernel
    # `__init__` parameter.
    del parameters['compute_acceptance']
    self._parameters = parameters

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

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

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

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

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

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

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

  @property
  def is_calibrated(self):
    return True

  def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of MALA.

    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 `step_size` or a list with same length as
        `current_state` or `diffusion_drift`.
    """
    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 UncalibratedLangevin(kernel_base.TransitionKernel):
  """Runs one step of Uncalibrated Langevin discretized diffusion.

  The class generates a Langevin proposal using `_euler_method` function and
  also computes helper `UncalibratedLangevinKernelResults` for the next
  iteration.

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

  For more details on `UncalibratedLangevin`, see
  `MetropolisAdjustedLangevinAlgorithm`.
  """

  def __init__(self,
               target_log_prob_fn,
               step_size,
               volatility_fn=None,
               parallel_iterations=10,
               compute_acceptance=True,
               seed=None,
               name=None):
    """Initializes Langevin diffusion 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.
      step_size: `Tensor` or Python `list` of `Tensor`s representing the step
        size for the leapfrog integrator. Must broadcast with the shape of
        `current_state`. Larger step sizes lead to faster progress, but
        too-large step sizes make rejection exponentially more likely. When
        possible, it's often helpful to match per-variable step sizes to the
        standard deviations of the target distribution in each variable.
      volatility_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns
        volatility value at `current_state`. Should return a `Tensor` or Python
        `list` of `Tensor`s that must broadcast with the shape of
        `current_state` Defaults to the identity function.
      parallel_iterations: the number of coordinates for which the gradients of
        the volatility matrix `volatility_fn` can be computed in parallel.
      compute_acceptance: Python 'bool' indicating whether to compute the
        Metropolis log-acceptance ratio used to construct
        `MetropolisAdjustedLangevinAlgorithm` kernel.
      seed: Python integer to seed the random number generator.
        Default value: `None` (i.e., no seed).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'mala_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 `step_size` or a list with same length as
        `current_state`.
      TypeError: if `volatility_fn` is not callable.
    """
    self._seed_stream = distributions.SeedStream(
        seed, salt='UncalibratedLangevin')
    # Default value of `volatility_fn` is the identity function.
    if volatility_fn is None:
      volatility_fn = lambda *args: 1.
    if not callable(volatility_fn):
      raise TypeError('`volatility_fn` must be callable (saw: {})'.format(
          type(volatility_fn)))
    self._parameters = dict(
        target_log_prob_fn=target_log_prob_fn,
        step_size=step_size,
        volatility_fn=volatility_fn,
        compute_acceptance=tf.convert_to_tensor(value=compute_acceptance),
        seed=seed,
        parallel_iterations=parallel_iterations,
        name=name)

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

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

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

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

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

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

  @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(MetropolisAdjustedLangevinAlgorithm.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, 'mala', 'one_step'),
        values=[
            self.step_size, current_state,
            previous_kernel_results.target_log_prob,
            previous_kernel_results.grads_target_log_prob,
            previous_kernel_results.volatility,
            previous_kernel_results.diffusion_drift
        ]):
      with tf.compat.v1.name_scope('initialize'):
        # Prepare input arguments to be passed to `_euler_method`.
        [
            current_state_parts,
            step_size_parts,
            current_target_log_prob,
            _,  # grads_target_log_prob
            current_volatility_parts,
            _,  # grads_volatility
            current_drift_parts,
        ] = _prepare_args(
            self.target_log_prob_fn,
            self.volatility_fn,
            current_state,
            self.step_size,
            previous_kernel_results.target_log_prob,
            previous_kernel_results.grads_target_log_prob,
            previous_kernel_results.volatility,
            previous_kernel_results.grads_volatility,
            previous_kernel_results.diffusion_drift,
            self.parallel_iterations)

        random_draw_parts = []
        for s in current_state_parts:
          random_draw_parts.append(
              tf.random.normal(
                  shape=tf.shape(input=s),
                  dtype=s.dtype.base_dtype,
                  seed=self._seed_stream()))

      # Number of independent chains run by the algorithm.
      independent_chain_ndims = distribution_util.prefer_static_rank(
          current_target_log_prob)

      # Generate the next state of the algorithm using Euler-Maruyama method.
      next_state_parts = _euler_method(random_draw_parts,
                                       current_state_parts,
                                       current_drift_parts,
                                       step_size_parts,
                                       current_volatility_parts)

      # Compute helper `UncalibratedLangevinKernelResults` to be processed by
      # `_compute_log_acceptance_correction` and in the next iteration of
      # `one_step` function.
      [
          _,  # state_parts
          _,  # step_sizes
          next_target_log_prob,
          next_grads_target_log_prob,
          next_volatility_parts,
          next_grads_volatility,
          next_drift_parts,
      ] = _prepare_args(
          self.target_log_prob_fn,
          self.volatility_fn,
          next_state_parts,
          step_size_parts,
          parallel_iterations=self.parallel_iterations)

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

      # Decide whether to compute the acceptance ratio
      log_acceptance_correction_compute = _compute_log_acceptance_correction(
          current_state_parts,
          next_state_parts,
          current_volatility_parts,
          next_volatility_parts,
          current_drift_parts,
          next_drift_parts,
          step_size_parts,
          independent_chain_ndims)
      log_acceptance_correction_skip = tf.zeros_like(next_target_log_prob)

      log_acceptance_correction = tf.cond(
          pred=self.compute_acceptance,
          true_fn=lambda: log_acceptance_correction_compute,
          false_fn=lambda: log_acceptance_correction_skip)

      return [
          maybe_flatten(next_state_parts),
          UncalibratedLangevinKernelResults(
              log_acceptance_correction=log_acceptance_correction,
              target_log_prob=next_target_log_prob,
              grads_target_log_prob=next_grads_target_log_prob,
              volatility=maybe_flatten(next_volatility_parts),
              grads_volatility=next_grads_volatility,
              diffusion_drift=next_drift_parts
          ),
      ]

  @mcmc_util.set_doc(
      MetropolisAdjustedLangevinAlgorithm.bootstrap_results.__doc__)
  def bootstrap_results(self, init_state):
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'mala', 'bootstrap_results'),
        values=[init_state]):
      init_state_parts = (list(init_state)
                          if mcmc_util.is_list_like(init_state)
                          else [init_state])

      init_state_parts = [
          tf.convert_to_tensor(value=x) for x in init_state_parts
      ]
      init_volatility = self.volatility_fn(*init_state_parts)  # pylint: disable=not-callable

      [
          _,  # state_parts
          _,  # step_sizes
          init_target_log_prob,
          init_grads_target_log_prob,
          init_volatility,
          init_grads_volatility,
          init_diffusion_drift,
      ] = _prepare_args(
          self.target_log_prob_fn,
          self.volatility_fn,
          state=init_state_parts,
          step_size=self.step_size,
          volatility=init_volatility,
          parallel_iterations=self.parallel_iterations)

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

      return UncalibratedLangevinKernelResults(
          log_acceptance_correction=tf.zeros_like(init_target_log_prob),
          target_log_prob=init_target_log_prob,
          grads_target_log_prob=init_grads_target_log_prob,
          volatility=maybe_flatten(init_volatility),
          grads_volatility=init_grads_volatility,
          diffusion_drift=init_diffusion_drift
      )

def _euler_method(random_draw_parts,
                  state_parts,
                  drift_parts,
                  step_size_parts,
                  volatility_parts,
                  name=None):
  """Applies one step of Euler-Maruyama method.

  Generates proposal of the form:

  ```python
  tfd.Normal(loc=state_parts + _get_drift(state_parts, ...),
             scale=tf.sqrt(step_size * volatility_fn(current_state)))
  ```

  `_get_drift(state_parts, ..)` is a diffusion drift value at `state_parts`.

  Args:
    random_draw_parts: Python `list` of `Tensor`s containing the value(s) of the
      random perturbation variable(s). Must broadcast with the shape of
      `state_parts`.
    state_parts: Python `list` of `Tensor`s representing the current
      state(s) of the Markov chain(s).
    drift_parts: Python `list` of `Tensor`s representing value of the drift
      `_get_drift(*state_parts, ..)`. Must broadcast with the shape of
      `state_parts`.
    step_size_parts: Python `list` of `Tensor`s representing the step size for
      the Euler-Maruyama method. Must broadcast with the shape of
      `state_parts`.  Larger step sizes lead to faster progress, but
      too-large step sizes make rejection exponentially more likely. When
      possible, it's often helpful to match per-variable step sizes to the
      standard deviations of the target distribution in each variable.
    volatility_parts: Python `list` of `Tensor`s representing the value of
      `volatility_fn(*state_parts)`. Must broadcast with the shape of
      `state_parts`.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'mala_euler_method').

  Returns:
    proposed_state_parts: Tensor or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at each result step. Has same shape as
      input `current_state_parts`.
  """
  with tf.compat.v1.name_scope(name, 'mala_euler_method', [
      random_draw_parts, state_parts, drift_parts, step_size_parts,
      volatility_parts
  ]):
    proposed_state_parts = []
    for random_draw, state, drift, step_size, volatility in zip(
        random_draw_parts,
        state_parts,
        drift_parts,
        step_size_parts,
        volatility_parts):
      proposal = state + drift + volatility * tf.sqrt(step_size) * random_draw
      proposed_state_parts.append(proposal)

    return proposed_state_parts

def _get_drift(step_size_parts, volatility_parts, grads_volatility,
               grads_target_log_prob,
               name=None):
  """Compute diffusion drift at the current location `current_state`.

  The drift of the diffusion at is computed as

  ```none
  0.5 * `step_size` * volatility_parts * `target_log_prob_fn(current_state)`
  + `step_size` * `grads_volatility`
  ```

  where `volatility_parts` = `volatility_fn(current_state)**2` and
  `grads_volatility` is a gradient of `volatility_parts` at the `current_state`.

  Args:
    step_size_parts: Python `list` of `Tensor`s representing the step size for
      Euler-Maruyama method. Must broadcast with the shape of
      `volatility_parts`.  Larger step sizes lead to faster progress, but
      too-large step sizes make rejection exponentially more likely. When
      possible, it's often helpful to match per-variable step sizes to the
      standard deviations of the target distribution in each variable.
    volatility_parts: Python `list` of `Tensor`s representing the value of
      `volatility_fn(*state_parts)`.
    grads_volatility: Python list of `Tensor`s representing the value of the
      gradient of `volatility_parts**2` wrt the state of the chain.
    grads_target_log_prob: Python list of `Tensor`s representing
      gradient of `target_log_prob_fn(*state_parts`) wrt `state_parts`. Must
      have same shape as `volatility_parts`.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'mala_get_drift').

  Returns:
    drift_parts: Tensor or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at each result step. Has same shape as
      input `current_state_parts`.
  """

  with tf.compat.v1.name_scope(name, 'mala_get_drift', [
      step_size_parts, volatility_parts, grads_volatility, grads_target_log_prob
  ]):

    drift_parts = []

    for step_size, volatility, grad_volatility, grad_target_log_prob in (
        zip(step_size_parts,
            volatility_parts,
            grads_volatility,
            grads_target_log_prob)):
      volatility_squared = tf.square(volatility)
      drift = 0.5 * step_size * (volatility_squared * grad_target_log_prob
                                 + grad_volatility)
      drift_parts.append(drift)

    return drift_parts

def _compute_log_acceptance_correction(current_state_parts,
                                       proposed_state_parts,
                                       current_volatility_parts,
                                       proposed_volatility_parts,
                                       current_drift_parts,
                                       proposed_drift_parts,
                                       step_size_parts,
                                       independent_chain_ndims,
                                       name=None):
  r"""Helper to `kernel` which computes the log acceptance-correction.

  Computes `log_acceptance_correction` as described in `MetropolisHastings`
  class. The proposal density is normal. More specifically,

   ```none
  q(proposed_state | current_state) \sim N(current_state + current_drift,
  step_size * current_volatility**2)

  q(current_state | proposed_state) \sim N(proposed_state + proposed_drift,
  step_size * proposed_volatility**2)
  ```

  The `log_acceptance_correction` is then

  ```none
  log_acceptance_correctio = q(current_state | proposed_state)
  - q(proposed_state | current_state)
  ```

  Args:
    current_state_parts: Python `list` of `Tensor`s representing the value(s) of
      the current state of the chain.
    proposed_state_parts:  Python `list` of `Tensor`s representing the value(s)
      of the proposed state of the chain. Must broadcast with the shape of
      `current_state_parts`.
    current_volatility_parts: Python `list` of `Tensor`s representing the value
      of `volatility_fn(*current_volatility_parts)`. Must broadcast with the
      shape of `current_state_parts`.
    proposed_volatility_parts: Python `list` of `Tensor`s representing the value
      of `volatility_fn(*proposed_volatility_parts)`. Must broadcast with the
      shape of `current_state_parts`
    current_drift_parts: Python `list` of `Tensor`s representing value of the
      drift `_get_drift(*current_state_parts, ..)`. Must broadcast with the
      shape of `current_state_parts`.
    proposed_drift_parts: Python `list` of `Tensor`s representing value of the
      drift `_get_drift(*proposed_drift_parts, ..)`. Must broadcast with the
      shape of `current_state_parts`.
    step_size_parts: Python `list` of `Tensor`s representing the step size for
      Euler-Maruyama method. Must broadcast with the shape of
      `current_state_parts`.
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'compute_log_acceptance_correction').

  Returns:
    log_acceptance_correction: `Tensor` representing the `log`
      acceptance-correction.  (See docstring for mathematical definition.)
  """

  with tf.compat.v1.name_scope(name, 'compute_log_acceptance_correction', [
      current_state_parts, proposed_state_parts, current_volatility_parts,
      proposed_volatility_parts, current_drift_parts, proposed_drift_parts,
      step_size_parts, independent_chain_ndims
  ]):

    proposed_log_density_parts = []
    dual_log_density_parts = []

    for [
        current_state,
        proposed_state,
        current_volatility,
        proposed_volatility,
        current_drift,
        proposed_drift,
        step_size,
    ] in zip(
        current_state_parts,
        proposed_state_parts,
        current_volatility_parts,
        proposed_volatility_parts,
        current_drift_parts,
        proposed_drift_parts,
        step_size_parts,
    ):
      axis = tf.range(independent_chain_ndims, tf.rank(current_state))

      state_diff = proposed_state - current_state

      current_volatility *= tf.sqrt(step_size)

      proposed_energy = (state_diff - current_drift) / current_volatility

      proposed_volatility *= tf.sqrt(step_size)
      # Compute part of `q(proposed_state | current_state)`
      proposed_energy = (
          tf.reduce_sum(
              input_tensor=mcmc_util.safe_sum(
                  [tf.math.log(current_volatility),
                   0.5 * (proposed_energy**2)]),
              axis=axis))
      proposed_log_density_parts.append(-proposed_energy)

      # Compute part of `q(current_state | proposed_state)`
      dual_energy = (state_diff + proposed_drift) / proposed_volatility
      dual_energy = (
          tf.reduce_sum(
              input_tensor=mcmc_util.safe_sum(
                  [tf.math.log(proposed_volatility), 0.5 * (dual_energy**2)]),
              axis=axis))
      dual_log_density_parts.append(-dual_energy)

    # Compute `q(proposed_state | current_state)`
    proposed_log_density_reduce = tf.reduce_sum(
        input_tensor=tf.stack(proposed_log_density_parts, axis=-1), axis=-1)
    # Compute `q(current_state | proposed_state)`
    dual_log_density_reduce = tf.reduce_sum(
        input_tensor=tf.stack(dual_log_density_parts, axis=-1), axis=-1)

    return mcmc_util.safe_sum([dual_log_density_reduce,
                               -proposed_log_density_reduce])

def _maybe_call_volatility_fn_and_grads(volatility_fn,
                                        state,
                                        volatility_fn_results=None,
                                        grads_volatility_fn=None,
                                        sample_shape=None,
                                        parallel_iterations=10):
  """Helper which computes `volatility_fn` results and grads, if needed."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  needs_volatility_fn_gradients = grads_volatility_fn is None

  # Convert `volatility_fn_results` to a list
  if volatility_fn_results is None:
    volatility_fn_results = volatility_fn(*state_parts)

  volatility_fn_results = (list(volatility_fn_results)
                           if mcmc_util.is_list_like(volatility_fn_results)
                           else [volatility_fn_results])
  if len(volatility_fn_results) == 1:
    volatility_fn_results *= len(state_parts)
  if len(state_parts) != len(volatility_fn_results):
    raise ValueError('`volatility_fn` should return a tensor or a list '
                     'of the same length as `current_state`.')

  # The shape of 'volatility_parts' needs to have the number of chains as a
  # leading dimension. For determinism we broadcast 'volatility_parts' to the
  # shape of `state_parts` since each dimension of `state_parts` could have a
  # different volatility value.

  volatility_fn_results = _maybe_broadcast_volatility(volatility_fn_results,
                                                      state_parts)
  if grads_volatility_fn is None:
    [
        _,
        grads_volatility_fn,
    ] = diag_jacobian(
        xs=state_parts,
        ys=volatility_fn_results,
        sample_shape=sample_shape,
        parallel_iterations=parallel_iterations,
        fn=volatility_fn)

  # Compute gradient of `volatility_parts**2`
  if needs_volatility_fn_gradients:
    grads_volatility_fn = [
        2. * g * volatility if g is not None else tf.zeros_like(
            fn_arg, dtype=fn_arg.dtype.base_dtype)
        for g, volatility, fn_arg in zip(
            grads_volatility_fn, volatility_fn_results, state_parts)
    ]

  return volatility_fn_results, grads_volatility_fn

def _maybe_broadcast_volatility(volatility_parts,
                                state_parts):
  """Helper to broadcast `volatility_parts` to the shape of `state_parts`."""
  return [v + tf.zeros_like(sp, dtype=sp.dtype.base_dtype)
          for v, sp in zip(volatility_parts, state_parts)]

def _prepare_args(target_log_prob_fn,
                  volatility_fn,
                  state,
                  step_size,
                  target_log_prob=None,
                  grads_target_log_prob=None,
                  volatility=None,
                  grads_volatility_fn=None,
                  diffusion_drift=None,
                  parallel_iterations=10):
  """Helper which processes input args to meet list-like assumptions."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]

  [
      target_log_prob,
      grads_target_log_prob,
  ] = mcmc_util.maybe_call_fn_and_grads(
      target_log_prob_fn,
      state_parts,
      target_log_prob,
      grads_target_log_prob)
  [
      volatility_parts,
      grads_volatility,
  ] = _maybe_call_volatility_fn_and_grads(
      volatility_fn,
      state_parts,
      volatility,
      grads_volatility_fn,
      distribution_util.prefer_static_shape(target_log_prob),
      parallel_iterations)

  step_sizes = (list(step_size) if mcmc_util.is_list_like(step_size)
                else [step_size])
  step_sizes = [
      tf.convert_to_tensor(
          value=s, name='step_size', dtype=target_log_prob.dtype)
      for s in step_sizes
  ]
  if len(step_sizes) == 1:
    step_sizes *= len(state_parts)
  if len(state_parts) != len(step_sizes):
    raise ValueError('There should be exactly one `step_size` or it should '
                     'have same length as `current_state`.')

  if diffusion_drift is None:
    diffusion_drift_parts = _get_drift(step_sizes, volatility_parts,
                                       grads_volatility,
                                       grads_target_log_prob)
  else:
    diffusion_drift_parts = (list(diffusion_drift)
                             if mcmc_util.is_list_like(diffusion_drift)
                             else [diffusion_drift])
    if len(state_parts) != len(diffusion_drift):
      raise ValueError('There should be exactly one `diffusion_drift` or it '
                       'should have same length as list-like `current_state`.')

  return [
      state_parts,
      step_sizes,
      target_log_prob,
      grads_target_log_prob,
      volatility_parts,
      grads_volatility,
      diffusion_drift_parts,
  ]