# 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.
# ============================================================================
# Functions "ndtr" and "ndtri" are derived from calculations made in:
# https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
# In the following email exchange, the author gives his consent to redistribute
# derived works under an Apache 2.0 license.
#
# From: Stephen Moshier <steve@moshier.net>
# Date: Sat, Jun 9, 2018 at 2:36 PM
# Subject: Re: Licensing cephes under Apache (BSD-like) license.
# To: rif <rif@google.com>
#
#
#
# Hello Rif,
#
# Yes, Google may distribute Cephes files under the Apache 2 license.
#
# If clarification is needed, I do not favor BSD over other free licenses.
# I would agree that Apache 2 seems to cover the concern you mentioned
# about sublicensees.
#
# Best wishes for good luck with your projects!
# Steve Moshier
#
#
#
# On Thu, 31 May 2018, rif wrote:
#
# > Hello Steve.
# > My name is Rif. I work on machine learning software at Google.
# >
# > Your cephes software continues to be incredibly useful and widely used. I
# > was wondering whether it would be permissible for us to use the Cephes code
# > under the Apache 2.0 license, which is extremely similar in permissions to
# > the BSD license (Wikipedia comparisons). This would be quite helpful to us
# > in terms of avoiding multiple licenses on software.
# >
# > I'm sorry to bother you with this (I can imagine you're sick of hearing
# > about this by now), but I want to be absolutely clear we're on the level and
# > not misusing your important software. In former conversation with Eugene
# > Brevdo (ebrevdo@google.com), you wrote "If your licensing is similar to BSD,
# > the formal way that has been handled is simply to add a statement to the
# > effect that you are incorporating the Cephes software by permission of the
# > author." I wanted to confirm that (a) we could use the Apache license, (b)
# > that we don't need to (and probably you don't want to) keep getting
# > contacted about individual uses, because your intent is generally to allow
# > this software to be reused under "BSD-like" license, and (c) you're OK
# > letting incorporators decide whether a license is sufficiently BSD-like?
# >
# > Best,
# >
# > rif
# >
# >
# >

"""Special Math Ops."""

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.internal import dtype_util

__all__ = [
    "erfinv",
    "ndtr",
    "ndtri",
    "log_ndtr",
    "log_cdf_laplace",
]

# log_ndtr uses different functions over the ranges
# (-infty, lower](lower, upper](upper, infty)
# Lower bound values were chosen by examining where the support of ndtr
# appears to be zero, relative to scipy's (which is always 64bit). They were
# then made more conservative just to be safe. (Conservative means use the
# expansion more than we probably need to.) See `NdtrTest` in
# special_math_test.py.
LOGNDTR_FLOAT64_LOWER = np.array(-20, np.float64)
LOGNDTR_FLOAT32_LOWER = np.array(-10, np.float32)

# Upper bound values were chosen by examining for which values of 'x'
# Log[cdf(x)] is 0, after which point we need to use the approximation
# Log[cdf(x)] = Log[1 - cdf(-x)] approx -cdf(-x). We chose a value slightly
# conservative, meaning we use the approximation earlier than needed.
LOGNDTR_FLOAT64_UPPER = np.array(8, np.float64)
LOGNDTR_FLOAT32_UPPER = np.array(5, np.float32)

def ndtr(x, name="ndtr"):
  """Normal distribution function.

  Returns the area under the Gaussian probability density function, integrated
  from minus infinity to x:

  ```
                    1       / x
     ndtr(x)  = ----------  |    exp(-0.5 t**2) dt
                sqrt(2 pi)  /-inf

              = 0.5 (1 + erf(x / sqrt(2)))
              = 0.5 erfc(x / sqrt(2))
  ```

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="ndtr").

  Returns:
    ndtr: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x` is not floating-type.
  """

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")
    if dtype_util.as_numpy_dtype(x.dtype) not in [np.float32, np.float64]:
      raise TypeError(
          "x.dtype=%s is not handled, see docstring for supported types."
          % x.dtype)
    return _ndtr(x)

def _ndtr(x):
  """Implements ndtr core logic."""
  half_sqrt_2 = tf.constant(
      0.5 * np.sqrt(2.), dtype=x.dtype, name="half_sqrt_2")
  w = x * half_sqrt_2
  z = tf.abs(w)
  y = tf.where(
      tf.less(z, half_sqrt_2), 1. + tf.math.erf(w),
      tf.where(tf.greater(w, 0.), 2. - tf.math.erfc(z), tf.math.erfc(z)))
  return 0.5 * y

def ndtri(p, name="ndtri"):
  """The inverse of the CDF of the Normal distribution function.

  Returns x such that the area under the pdf from minus infinity to x is equal
  to p.

  A piece-wise rational approximation is done for the function.
  This is a port of the implementation in netlib.

  Args:
    p: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="ndtri").

  Returns:
    x: `Tensor` with `dtype=p.dtype`.

  Raises:
    TypeError: if `p` is not floating-type.
  """

  with tf.name_scope(name):
    p = tf.convert_to_tensor(value=p, name="p")
    if dtype_util.as_numpy_dtype(p.dtype) not in [np.float32, np.float64]:
      raise TypeError(
          "p.dtype=%s is not handled, see docstring for supported types."
          % p.dtype)
    return _ndtri(p)

def _ndtri(p):
  """Implements ndtri core logic."""

  # Constants used in piece-wise rational approximations. Taken from the cephes
  # library:
  # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
  p0 = list(reversed([-5.99633501014107895267E1,
                      9.80010754185999661536E1,
                      -5.66762857469070293439E1,
                      1.39312609387279679503E1,
                      -1.23916583867381258016E0]))
  q0 = list(reversed([1.0,
                      1.95448858338141759834E0,
                      4.67627912898881538453E0,
                      8.63602421390890590575E1,
                      -2.25462687854119370527E2,
                      2.00260212380060660359E2,
                      -8.20372256168333339912E1,
                      1.59056225126211695515E1,
                      -1.18331621121330003142E0]))
  p1 = list(reversed([4.05544892305962419923E0,
                      3.15251094599893866154E1,
                      5.71628192246421288162E1,
                      4.40805073893200834700E1,
                      1.46849561928858024014E1,
                      2.18663306850790267539E0,
                      -1.40256079171354495875E-1,
                      -3.50424626827848203418E-2,
                      -8.57456785154685413611E-4]))
  q1 = list(reversed([1.0,
                      1.57799883256466749731E1,
                      4.53907635128879210584E1,
                      4.13172038254672030440E1,
                      1.50425385692907503408E1,
                      2.50464946208309415979E0,
                      -1.42182922854787788574E-1,
                      -3.80806407691578277194E-2,
                      -9.33259480895457427372E-4]))
  p2 = list(reversed([3.23774891776946035970E0,
                      6.91522889068984211695E0,
                      3.93881025292474443415E0,
                      1.33303460815807542389E0,
                      2.01485389549179081538E-1,
                      1.23716634817820021358E-2,
                      3.01581553508235416007E-4,
                      2.65806974686737550832E-6,
                      6.23974539184983293730E-9]))
  q2 = list(reversed([1.0,
                      6.02427039364742014255E0,
                      3.67983563856160859403E0,
                      1.37702099489081330271E0,
                      2.16236993594496635890E-1,
                      1.34204006088543189037E-2,
                      3.28014464682127739104E-4,
                      2.89247864745380683936E-6,
                      6.79019408009981274425E-9]))

  def _create_polynomial(var, coeffs):
    """Compute n_th order polynomial via Horner's method."""
    coeffs = np.array(coeffs, dtype_util.as_numpy_dtype(var.dtype))
    if not coeffs.size:
      return tf.zeros_like(var)
    return coeffs[0] + _create_polynomial(var, coeffs[1:]) * var

  maybe_complement_p = tf.where(p > -np.expm1(-2.), 1. - p, p)
  # Write in an arbitrary value in place of 0 for p since 0 will cause NaNs
  # later on. The result from the computation when p == 0 is not used so any
  # number that doesn't result in NaNs is fine.
  sanitized_mcp = tf.where(
      maybe_complement_p <= 0.,
      tf.fill(tf.shape(input=p),
              dtype_util.as_numpy_dtype(p.dtype)(0.5)), maybe_complement_p)

  # Compute x for p > exp(-2): x/sqrt(2pi) = w + w**3 P0(w**2)/Q0(w**2).
  w = sanitized_mcp - 0.5
  ww = w ** 2
  x_for_big_p = w + w * ww * (_create_polynomial(ww, p0)
                              / _create_polynomial(ww, q0))
  x_for_big_p *= -np.sqrt(2. * np.pi)

  # Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
  # where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
  # arrays based on whether p < exp(-32).
  z = tf.sqrt(-2. * tf.math.log(sanitized_mcp))
  first_term = z - tf.math.log(z) / z
  second_term_small_p = (
      _create_polynomial(1. / z, p2) /
      _create_polynomial(1. / z, q2) / z)
  second_term_otherwise = (
      _create_polynomial(1. / z, p1) /
      _create_polynomial(1. / z, q1) / z)
  x_for_small_p = first_term - second_term_small_p
  x_otherwise = first_term - second_term_otherwise

  x = tf.where(sanitized_mcp > np.exp(-2.), x_for_big_p,
               tf.where(z >= 8.0, x_for_small_p, x_otherwise))

  x = tf.where(p > 1. - np.exp(-2.), x, -x)
  infinity_scalar = tf.constant(np.inf, dtype=p.dtype)
  infinity = tf.fill(tf.shape(input=p), infinity_scalar)
  x_nan_replaced = tf.where(
      p <= 0.0, -infinity, tf.where(p >= 1.0, infinity, x))
  return x_nan_replaced

def log_ndtr(x, series_order=3, name="log_ndtr"):
  """Log Normal distribution function.

  For details of the Normal distribution function see `ndtr`.

  This function calculates `(log o ndtr)(x)` by either calling `log(ndtr(x))` or
  using an asymptotic series. Specifically:
  - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
    `log(1-x) ~= -x, x << 1`.
  - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique
    and take a log.
  - For `x <= lower_segment`, we use the series approximation of erf to compute
    the log CDF directly.

  The `lower_segment` is set based on the precision of the input:

  ```
  lower_segment = { -20,  x.dtype=float64
                  { -10,  x.dtype=float32
  upper_segment = {   8,  x.dtype=float64
                  {   5,  x.dtype=float32
  ```

  When `x < lower_segment`, the `ndtr` asymptotic series approximation is:

  ```
     ndtr(x) = scale * (1 + sum) + R_N
     scale   = exp(-0.5 x**2) / (-x sqrt(2 pi))
     sum     = Sum{(-1)^n (2n-1)!! / (x**2)^n, n=1:N}
     R_N     = O(exp(-0.5 x**2) (2N+1)!! / |x|^{2N+3})
  ```

  where `(2n-1)!! = (2n-1) (2n-3) (2n-5) ...  (3) (1)` is a
  [double-factorial](https://en.wikipedia.org/wiki/Double_factorial).

  Args:
    x: `Tensor` of type `float32`, `float64`.
    series_order: Positive Python `integer`. Maximum depth to
      evaluate the asymptotic expansion. This is the `N` above.
    name: Python string. A name for the operation (default="log_ndtr").

  Returns:
    log_ndtr: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x.dtype` is not handled.
    TypeError: if `series_order` is a not Python `integer.`
    ValueError:  if `series_order` is not in `[0, 30]`.
  """
  if not isinstance(series_order, int):
    raise TypeError("series_order must be a Python integer.")
  if series_order < 0:
    raise ValueError("series_order must be non-negative.")
  if series_order > 30:
    raise ValueError("series_order must be <= 30.")

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")

    if dtype_util.base_equal(x.dtype, tf.float64):
      lower_segment = LOGNDTR_FLOAT64_LOWER
      upper_segment = LOGNDTR_FLOAT64_UPPER
    elif dtype_util.base_equal(x.dtype, tf.float32):
      lower_segment = LOGNDTR_FLOAT32_LOWER
      upper_segment = LOGNDTR_FLOAT32_UPPER
    else:
      raise TypeError("x.dtype=%s is not supported." % x.dtype)

    # The basic idea here was ported from:
    #   https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
    # We copy the main idea, with a few changes
    # * For x >> 1, and X ~ Normal(0, 1),
    #     Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
    #     which extends the range of validity of this function.
    # * We use one fixed series_order for all of 'x', rather than adaptive.
    # * Our docstring properly reflects that this is an asymptotic series, not a
    #   Taylor series. We also provided a correct bound on the remainder.
    # * We need to use the max/min in the _log_ndtr_lower arg to avoid nan when
    #   x=0. This happens even though the branch is unchosen because when x=0
    #   the gradient of a select involves the calculation 1*dy+0*(-inf)=nan
    #   regardless of whether dy is finite. Note that the minimum is a NOP if
    #   the branch is chosen.
    return tf.where(
        tf.greater(x, upper_segment),
        -_ndtr(-x),  # log(1-x) ~= -x, x << 1
        tf.where(
            tf.greater(x, lower_segment),
            tf.math.log(_ndtr(tf.maximum(x, lower_segment))),
            _log_ndtr_lower(tf.minimum(x, lower_segment), series_order)))

def _log_ndtr_lower(x, series_order):
  """Asymptotic expansion version of `Log[cdf(x)]`, appropriate for `x<<-1`."""
  x_2 = tf.square(x)
  # Log of the term multiplying (1 + sum)
  log_scale = -0.5 * x_2 - tf.math.log(-x) - 0.5 * np.log(2. * np.pi)
  return log_scale + tf.math.log(_log_ndtr_asymptotic_series(x, series_order))

def _log_ndtr_asymptotic_series(x, series_order):
  """Calculates the asymptotic series used in log_ndtr."""
  npdt = dtype_util.as_numpy_dtype(x.dtype)
  if series_order <= 0:
    return npdt(1)
  x_2 = tf.square(x)
  even_sum = tf.zeros_like(x)
  odd_sum = tf.zeros_like(x)
  x_2n = x_2  # Start with x^{2*1} = x^{2*n} with n = 1.
  for n in range(1, series_order + 1):
    y = npdt(_double_factorial(2 * n - 1)) / x_2n
    if n % 2:
      odd_sum += y
    else:
      even_sum += y
    x_2n *= x_2
  return 1. + even_sum - odd_sum

def erfinv(x, name="erfinv"):
  """The inverse function for erf, the error function.

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="erfinv").

  Returns:
    x: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x` is not floating-type.
  """

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")
    if dtype_util.as_numpy_dtype(x.dtype) not in [np.float32, np.float64]:
      raise TypeError("x.dtype={} is not handled, see docstring for supported "
                      "types.".format(dtype_util.name(x.dtype)))
    return ndtri((x + 1.) / 2.) / np.sqrt(2.)

def _double_factorial(n):
  """The double factorial function for small Python integer `n`."""
  return np.prod(np.arange(n, 1, -2))

def log_cdf_laplace(x, name="log_cdf_laplace"):
  """Log Laplace distribution function.

  This function calculates `Log[L(x)]`, where `L(x)` is the cumulative
  distribution function of the Laplace distribution, i.e.

  ```L(x) := 0.5 * int_{-infty}^x e^{-|t|} dt```

  For numerical accuracy, `L(x)` is computed in different ways depending on `x`,

  ```
  x <= 0:
    Log[L(x)] = Log[0.5] + x, which is exact

  0 < x:
    Log[L(x)] = Log[1 - 0.5 * e^{-x}], which is exact
  ```

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="log_ndtr").

  Returns:
    `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x.dtype` is not handled.
  """

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")

    # For x < 0, L(x) = 0.5 * exp{x} exactly, so Log[L(x)] = log(0.5) + x.
    lower_solution = -np.log(2.) + x

    # safe_exp_neg_x = exp{-x} for x > 0, but is
    # bounded above by 1, which avoids
    #   log[1 - 1] = -inf for x = log(1/2), AND
    #   exp{-x} --> inf, for x << -1
    safe_exp_neg_x = tf.exp(-tf.abs(x))

    # log1p(z) = log(1 + z) approx z for |z| << 1. This approxmation is used
    # internally by log1p, rather than being done explicitly here.
    upper_solution = tf.math.log1p(-0.5 * safe_exp_neg_x)

    return tf.where(x < 0., lower_solution, upper_solution)


# 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.
# ============================================================================
"""Shared library for `text_messages_hmc_{graph,eager}_test.py`."""

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

import time
# Dependency imports
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions

__all__ = [
    'TextMessagesHmcBenchmarkTestHarness',
    'benchmark_text_messages_hmc',
    'text_messages_joint_log_prob',
]

def mvn(*args, **kwargs):
  """Convenience function to efficiently construct a MultivariateNormalDiag."""
  # Faster than using `tfd.MultivariateNormalDiag`.
  return tfd.Independent(tfd.Normal(*args, **kwargs),
                         reinterpreted_batch_ndims=1)

def text_messages_joint_log_prob(count_data, lambda_1, lambda_2, tau):
  """Joint log probability function."""
  alpha = (1. / tf.reduce_mean(input_tensor=count_data))
  rv_lambda = tfd.Exponential(rate=alpha)

  rv_tau = tfd.Uniform()

  lambda_ = tf.gather(
      [lambda_1, lambda_2],
      indices=tf.cast(
          tau * tf.cast(tf.size(input=count_data), dtype=tf.float32) <= tf.cast(
              tf.range(tf.size(input=count_data)), dtype=tf.float32),
          dtype=tf.int32))
  rv_observation = tfd.Poisson(rate=lambda_)

  return (rv_lambda.log_prob(lambda_1) + rv_lambda.log_prob(lambda_2) +
          rv_tau.log_prob(tau) +
          tf.reduce_sum(input_tensor=rv_observation.log_prob(count_data)))

def benchmark_text_messages_hmc(
    num_results=int(3e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3):
  """Runs HMC on the text-messages unnormalized posterior."""

  if not tf.executing_eagerly():
    tf.compat.v1.reset_default_graph()

  # Build a static, pretend dataset.
  count_data = tf.cast(
      tf.concat(
          [tfd.Poisson(rate=15.).sample(43),
           tfd.Poisson(rate=25.).sample(31)],
          axis=0),
      dtype=tf.float32)
  if tf.executing_eagerly():
    count_data = count_data.numpy()
  else:
    with tf.compat.v1.Session():
      count_data = count_data.eval()

  # Define a closure over our joint_log_prob.
  def unnormalized_log_posterior(lambda1, lambda2, tau):
    return text_messages_joint_log_prob(count_data, lambda1, lambda2, tau)

  if tf.executing_eagerly():
    sample_chain = tf.function(tfp.mcmc.sample_chain)
  else:
    sample_chain = tfp.mcmc.sample_chain

  # Initialize the step_size. (It will be automatically adapted.)
  step_size = tf.compat.v2.Variable(
      name='step_size',
      initial_value=tf.constant(0.05, dtype=tf.float32),
      trainable=False)

  def computation():
    """The benchmark computation."""

    initial_chain_state = [
        tf.constant(count_data.mean(), name='init_lambda1'),
        tf.constant(count_data.mean(), name='init_lambda2'),
        tf.constant(0.5, name='init_tau'),
    ]

    unconstraining_bijectors = [
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Sigmoid(),   # Maps [0,1] to R.
    ]

    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=initial_chain_state,
        kernel=tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=unnormalized_log_posterior,
                num_leapfrog_steps=num_leapfrog_steps,
                step_size=step_size,
                step_size_update_fn=
                tfp.mcmc.make_simple_step_size_update_policy(num_burnin_steps),
                state_gradients_are_stopped=True),
            bijector=unconstraining_bijectors))

    return kernel_results.inner_results.is_accepted

  # Let's force evaluation of graph to ensure build time is not part of our time
  # trial.
  is_accepted_tensor = computation()
  if not tf.executing_eagerly():
    session = tf.compat.v1.Session()
    session.run(tf.compat.v1.global_variables_initializer())
    session.run(is_accepted_tensor)

  start_time = time.time()
  if tf.executing_eagerly():
    is_accepted = computation()
  else:
    is_accepted = session.run(is_accepted_tensor)
  wall_time = time.time() - start_time

  num_accepted = np.sum(is_accepted)
  acceptance_rate = np.float32(num_accepted) / np.float32(num_results)

  return dict(
      iters=(num_results + num_burnin_steps) * num_leapfrog_steps,
      extras={'acceptance_rate': acceptance_rate},
      wall_time=wall_time)

class TextMessagesHmcBenchmarkTestHarness(object):
  """Test harness for running HMC benchmark tests in graph/eager modes."""

  def __init__(self):
    self._mode = 'eager' if tf.executing_eagerly() else 'graph'

  def benchmark_text_messages_hmc_num_leapfrog_1(self):
    self.report_benchmark(
        name=self._mode + '_text_messages_hmc_num_leapfrog_1',
        **benchmark_text_messages_hmc(num_leapfrog_steps=1))

  def benchmark_text_messages_hmc_num_leapfrog_2(self):
    self.report_benchmark(
        name=self._mode + '_text_messages_hmc_num_leapfrog_2',
        **benchmark_text_messages_hmc(num_leapfrog_steps=2))

  def benchmark_text_messages_hmc_num_leapfrog_3(self):
    self.report_benchmark(
        name=self._mode + '_text_messages_hmc_num_leapfrog_3',
        **benchmark_text_messages_hmc(num_leapfrog_steps=3))

  def benchmark_text_messages_hmc_num_leapfrog_10(self):
    self.report_benchmark(
        name=self._mode + '_text_messages_hmc_num_leapfrog_10',
        **benchmark_text_messages_hmc(num_leapfrog_steps=10))

  def benchmark_text_messages_hmc_num_leapfrog_20(self):
    self.report_benchmark(
        name=self._mode + '_text_messages_hmc_num_leapfrog_20',
        **benchmark_text_messages_hmc(num_leapfrog_steps=20))


# 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 GaussianProcess distribution class."""

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

import warnings

# Dependency imports
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.distributions import mvn_linear_operator
from tensorflow_probability.python.distributions import normal
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization

__all__ = [
    'GaussianProcess',
]

def _add_diagonal_shift(matrix, shift):
  return tf.linalg.set_diag(
      matrix, tf.linalg.diag_part(matrix) + shift, name='add_diagonal_shift')

class GaussianProcess(distribution.Distribution):
  """Marginal distribution of a Gaussian process at finitely many points.

  A Gaussian process (GP) is an indexed collection of random variables, any
  finite collection of which are jointly Gaussian. While this definition applies
  to finite index sets, it is typically implicit that the index set is infinite;
  in applications, it is often some finite dimensional real or complex vector
  space. In such cases, the GP may be thought of as a distribution over
  (real- or complex-valued) functions defined over the index set.

  Just as Gaussian distributions are fully specified by their first and second
  moments, a Gaussian process can be completely specified by a mean and
  covariance function. Let `S` denote the index set and `K` the space in which
  each indexed random variable takes its values (again, often R or C). The mean
  function is then a map `m: S -> K`, and the covariance function, or kernel, is
  a positive-definite function `k: (S x S) -> K`. The properties of functions
  drawn from a GP are entirely dictated (up to translation) by the form of the
  kernel function.

  This `Distribution` represents the marginal joint distribution over function
  values at a given finite collection of points `[x[1], ..., x[N]]` from the
  index set `S`. By definition, this marginal distribution is just a
  multivariate normal distribution, whose mean is given by the vector
  `[ m(x[1]), ..., m(x[N]) ]` and whose covariance matrix is constructed from
  pairwise applications of the kernel function to the given inputs:

  ```none
      | k(x[1], x[1])    k(x[1], x[2])  ...  k(x[1], x[N]) |
      | k(x[2], x[1])    k(x[2], x[2])  ...  k(x[2], x[N]) |
      |      ...              ...                 ...      |
      | k(x[N], x[1])    k(x[N], x[2])  ...  k(x[N], x[N]) |
  ```

  For this to be a valid covariance matrix, it must be symmetric and positive
  definite; hence the requirement that `k` be a positive definite function
  (which, by definition, says that the above procedure will yield PD matrices).

  We also support the inclusion of zero-mean Gaussian noise in the model, via
  the `observation_noise_variance` parameter. This augments the generative model
  to

  ```none
  f ~ GP(m, k)
  (y[i] | f, x[i]) ~ Normal(f(x[i]), s)
  ```

  where

    * `m` is the mean function
    * `k` is the covariance kernel function
    * `f` is the function drawn from the GP
    * `x[i]` are the index points at which the function is observed
    * `y[i]` are the observed values at the index points
    * `s` is the scale of the observation noise.

  Note that this class represents an *unconditional* Gaussian process; it does
  not implement posterior inference conditional on observed function
  evaluations. This class is useful, for example, if one wishes to combine a GP
  prior with a non-conjugate likelihood using MCMC to sample from the posterior.

  #### Mathematical Details

  The probability density function (pdf) is a multivariate normal whose
  parameters are derived from the GP's properties:

  ```none
  pdf(x; index_points, mean_fn, kernel) = exp(-0.5 * y) / Z
  K = (kernel.matrix(index_points, index_points) +
       (observation_noise_variance + jitter) * eye(N))
  y = (x - mean_fn(index_points))^T @ K @ (x - mean_fn(index_points))
  Z = (2 * pi)**(.5 * N) |det(K)|**(.5)
  ```

  where:

  * `index_points` are points in the index set over which the GP is defined,
  * `mean_fn` is a callable mapping the index set to the GP's mean values,
  * `kernel` is `PositiveSemidefiniteKernel`-like and represents the covariance
    function of the GP,
  * `observation_noise_variance` represents (optional) observation noise.
  * `jitter` is added to the diagonal to ensure positive definiteness up to
     machine precision (otherwise Cholesky-decomposition is prone to failure),
  * `eye(N)` is an N-by-N identity matrix.

  #### Examples

  ##### Draw joint samples from a GP prior

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

  tfd = tfp.distributions
  psd_kernels = tfp.positive_semidefinite_kernels

  num_points = 100
  # Index points should be a collection (100, here) of feature vectors. In this
  # example, we're using 1-d vectors, so we just need to reshape the output from
  # np.linspace, to give a shape of (100, 1).
  index_points = np.expand_dims(np.linspace(-1., 1., num_points), -1)

  # Define a kernel with default parameters.
  kernel = psd_kernels.ExponentiatedQuadratic()

  gp = tfd.GaussianProcess(kernel, index_points)

  samples = gp.sample(10)
  # ==> 10 independently drawn, joint samples at `index_points`

  noisy_gp = tfd.GaussianProcess(
      kernel=kernel,
      index_points=index_points,
      observation_noise_variance=.05)
  noisy_samples = noisy_gp.sample(10)
  # ==> 10 independently drawn, noisy joint samples at `index_points`
  ```

  ##### Optimize kernel parameters via maximum marginal likelihood.

  ```python
  # Suppose we have some data from a known function. Note the index points in
  # general have shape `[b1, ..., bB, f1, ..., fF]` (here we assume `F == 1`),
  # so we need to explicitly consume the feature dimensions (just the last one
  # here).
  f = lambda x: np.sin(10*x[..., 0]) * np.exp(-x[..., 0]**2)
  observed_index_points = np.expand_dims(np.random.uniform(-1., 1., 50), -1)
  # Squeeze to take the shape from [50, 1] to [50].
  observed_values = f(observed_index_points)

  # Define a kernel with trainable parameters.
  kernel = psd_kernels.ExponentiatedQuadratic(
      amplitude=tf.get_variable('amplitude', shape=(), dtype=np.float64),
      length_scale=tf.get_variable('length_scale', shape=(), dtype=np.float64))

  gp = tfd.GaussianProcess(kernel, observed_index_points)
  neg_log_likelihood = -gp.log_prob(observed_values)

  optimize = tf.train.AdamOptimizer().minimize(neg_log_likelihood)

  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    for i in range(1000):
      _, neg_log_likelihood_ = sess.run([optimize, neg_log_likelihood])
      if i % 100 == 0:
        print("Step {}: NLL = {}".format(i, neg_log_likelihood_))
    print("Final NLL = {}".format(neg_log_likelihood_))
  ```

  """

  def __init__(self,
               kernel,
               index_points=None,
               mean_fn=None,
               observation_noise_variance=0.,
               jitter=1e-6,
               validate_args=False,
               allow_nan_stats=False,
               name='GaussianProcess'):
    """Instantiate a GaussianProcess Distribution.

    Args:
      kernel: `PositiveSemidefiniteKernel`-like instance representing the
        GP's covariance function.
      index_points: `float` `Tensor` representing finite (batch of) vector(s) of
        points in the index set over which the GP is defined. Shape has the form
        `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e` is the number
        (size) of index points in each batch. Ultimately this distribution
        corresponds to a `e`-dimensional multivariate normal. The batch shape
        must be broadcastable with `kernel.batch_shape` and any batch dims
        yielded by `mean_fn`.
      mean_fn: Python `callable` that acts on `index_points` to produce a (batch
        of) vector(s) of mean values at `index_points`. Takes a `Tensor` of
        shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
        broadcastable with `[b1, ..., bB]`. Default value: `None` implies
        constant zero function.
      observation_noise_variance: `float` `Tensor` representing the variance
        of the noise in the Normal likelihood distribution of the model. May be
        batched, in which case the batch shape must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc.).
        Default value: `0.`
      jitter: `float` scalar `Tensor` added to the diagonal of the covariance
        matrix to ensure positive definiteness of the covariance matrix.
        Default value: `1e-6`.
      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.
        Default value: `False`.
      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.
        Default value: `False`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".

    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype(
          [index_points, observation_noise_variance, jitter], tf.float32)
      if index_points is not None:
        index_points = tf.convert_to_tensor(
            value=index_points, dtype=dtype, name='index_points')
      jitter = tf.convert_to_tensor(value=jitter, dtype=dtype, name='jitter')
      observation_noise_variance = tf.convert_to_tensor(
          value=observation_noise_variance,
          dtype=dtype,
          name='observation_noise_variance')

      self._kernel = kernel
      self._index_points = index_points
      # Default to a constant zero function, borrowing the dtype from
      # index_points to ensure consistency.
      if mean_fn is None:
        mean_fn = lambda x: tf.zeros([1], dtype=dtype)
      else:
        if not callable(mean_fn):
          raise ValueError('`mean_fn` must be a Python callable')
      self._mean_fn = mean_fn
      self._observation_noise_variance = observation_noise_variance
      self._jitter = jitter

      graph_parents = [observation_noise_variance, jitter]
      if index_points is not None: graph_parents.append(index_points)

      with tf.name_scope('init'):
        super(GaussianProcess, self).__init__(
            dtype=dtype,
            reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            parameters=parameters,
            graph_parents=graph_parents,
            name=name)

  def _is_univariate_marginal(self, index_points):
    """True if the given index_points would yield a univariate marginal.

    Args:
      index_points: the set of index set locations at which to compute the
      marginal Gaussian distribution. If this set is of size 1, the marginal is
      univariate.

    Returns:
      is_univariate: Boolean indicating whether the marginal is univariate or
      multivariate. In the case of dynamic shape in the number of index points,
      defaults to "multivariate" since that's the best we can do.
    """
    num_index_points = tf.compat.dimension_value(
        index_points.shape[-(self.kernel.feature_ndims + 1)])
    if num_index_points is None:
      warnings.warn(
          'Unable to detect statically whether the number of index_points is '
          '1. As a result, defaulting to treating the marginal GP at '
          '`index_points` as a multivariate Gaussian. This makes some methods, '
          'like `cdf` unavailable.')
    return num_index_points == 1

  def _compute_covariance(self, index_points):
    kernel_matrix = self.kernel.matrix(index_points, index_points)
    if self._is_univariate_marginal(index_points):
      # kernel_matrix thus has shape [..., 1, 1]; squeeze off the last dims and
      # tack on the observation noise variance.
      return (tf.squeeze(kernel_matrix, axis=[-2, -1]) +
              self.observation_noise_variance)
    else:
      return _add_diagonal_shift(
          kernel_matrix, self.observation_noise_variance)

  def get_marginal_distribution(self, index_points=None):
    """Compute the marginal of this GP over function values at `index_points`.

    Args:
      index_points: `float` `Tensor` representing finite (batch of) vector(s) of
        points in the index set over which the GP is defined. Shape has the form
        `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e` is the number
        (size) of index points in each batch. Ultimately this distribution
        corresponds to a `e`-dimensional multivariate normal. The batch shape
        must be broadcastable with `kernel.batch_shape` and any batch dims
        yielded by `mean_fn`.

    Returns:
      marginal: a `Normal` or `MultivariateNormalLinearOperator` distribution,
        according to whether `index_points` consists of one or many index
        points, respectively.
    """
    with self._name_scope('get_marginal_distribution'):
      # TODO(cgs): consider caching the result here, keyed on `index_points`.
      index_points = self._get_index_points(index_points)
      covariance = self._compute_covariance(index_points)
      loc = self._mean_fn(index_points)
      # If we're sure the number of index points is 1, we can just construct a
      # scalar Normal. This has computational benefits and supports things like
      # CDF that aren't otherwise straightforward to provide.
      if self._is_univariate_marginal(index_points):
        scale = tf.sqrt(covariance)
        # `loc` has a trailing 1 in the shape; squeeze it.
        loc = tf.squeeze(loc, axis=-1)
        return normal.Normal(
            loc=loc,
            scale=scale,
            validate_args=self._validate_args,
            allow_nan_stats=self._allow_nan_stats,
            name='marginal_distribution')
      else:
        scale = tf.linalg.LinearOperatorLowerTriangular(
            tf.linalg.cholesky(_add_diagonal_shift(covariance, self.jitter)),
            is_non_singular=True,
            name='GaussianProcessScaleLinearOperator')
        return mvn_linear_operator.MultivariateNormalLinearOperator(
            loc=loc,
            scale=scale,
            validate_args=self._validate_args,
            allow_nan_stats=self._allow_nan_stats,
            name='marginal_distribution')

  @property
  def mean_fn(self):
    return self._mean_fn

  @property
  def kernel(self):
    return self._kernel

  @property
  def index_points(self):
    return self._index_points

  @property
  def observation_noise_variance(self):
    return self._observation_noise_variance

  @property
  def jitter(self):
    return self._jitter

  def _get_index_points(self, index_points=None):
    """Return `index_points` if not None, else `self._index_points`.

    Args:
      index_points: if given, this is what is returned; else,
      `self._index_points`

    Returns:
      index_points: the given arg, if not None, else the class member
      `self._index_points`.

    Rases:
      ValueError: if `index_points` and `self._index_points` are both `None`.
    """
    if self._index_points is None and index_points is None:
      raise ValueError(
          'This GaussianProcess instance was not instantiated with a value for '
          'index_points. One must therefore be provided when calling sample, '
          'log_prob, and other such methods. In particular, one can\'t compute '
          'KL divergences to/from an instance of `GaussianProccess` with '
          'unspecified `index_points` directly. Instead, use the '
          '`get_marginal_distribution` function, which takes `index_points` as '
          'an argument and returns a `Normal` or '
          '`MultivariateNormalLinearOperator` instance, whose KL can be '
          'computed.')
    return index_points if index_points is not None else self._index_points

  def _log_prob(self, value, index_points=None):
    return self.get_marginal_distribution(index_points).log_prob(value)

  def _batch_shape_tensor(self, index_points=None):
    return self.get_marginal_distribution(index_points).batch_shape_tensor()

  def _batch_shape(self, index_points=None):
    return self.get_marginal_distribution(index_points).batch_shape

  def _event_shape_tensor(self, index_points=None):
    return self.get_marginal_distribution(index_points).event_shape_tensor()

  def _event_shape(self, index_points=None):
    return self.get_marginal_distribution(index_points).event_shape

  def _sample_n(self, n, seed=None, index_points=None):
    return self.get_marginal_distribution(index_points).sample(n, seed=seed)

  def _log_survival_function(self, value, index_points=None):
    return self.get_marginal_distribution(
        index_points).log_survival_function(value)

  def _survival_function(self, value, index_points=None):
    return self.get_marginal_distribution(index_points).survival_function(value)

  def _log_cdf(self, value, index_points=None):
    return self.get_marginal_distribution(index_points).log_cdf(value)

  def _entropy(self, index_points=None):
    return self.get_marginal_distribution(index_points).entropy()

  def _mean(self, index_points=None):
    return self.get_marginal_distribution(index_points).mean()

  def _quantile(self, value, index_points=None):
    return self.get_marginal_distribution(index_points).quantile(value)

  def _stddev(self, index_points=None):
    return tf.sqrt(self._variance(index_points=index_points))

  def _variance(self, index_points=None):
    index_points = self._get_index_points(index_points)

    kernel_diag = self.kernel.apply(index_points, index_points)
    if self._is_univariate_marginal(index_points):
      return (tf.squeeze(kernel_diag, axis=[-1]) +
              self.observation_noise_variance)
    else:
      return kernel_diag + self.observation_noise_variance

  def _covariance(self, index_points=None):
    # Using the result of get_marginal_distribution would involve an extra
    # matmul, and possibly even an unneceesary cholesky first. We can avoid that
    # by going straight through the kernel function.
    return self._compute_covariance(self._get_index_points(index_points))

  def _mode(self, index_points=None):
    return self.get_marginal_distribution(index_points).mode()

def _assert_kl_compatible(marginal, other):
  if ((isinstance(marginal, normal.Normal) and
       isinstance(other, normal.Normal)) or
      (isinstance(marginal,
                  mvn_linear_operator.MultivariateNormalLinearOperator) and
       isinstance(other,
                  mvn_linear_operator.MultivariateNormalLinearOperator))):
    return
  raise ValueError(
      'Attempting to compute KL between a GP marginal and a distribution of '
      'incompatible type. GP marginal has type {} and other distribution has '
      'type {}.'.format(type(marginal), type(other)))

def _kl_gp_compatible(gp, compatible, name):
  with tf.name_scope(name):
    marginal = gp.get_marginal_distribution()
    _assert_kl_compatible(marginal, compatible)
    return kullback_leibler.kl_divergence(marginal, compatible)

def _kl_compatible_gp(compatible, gp, name):
  with tf.name_scope(name):
    marginal = gp.get_marginal_distribution()
    _assert_kl_compatible(marginal, compatible)
    return kullback_leibler.kl_divergence(compatible, marginal)

@kullback_leibler.RegisterKL(GaussianProcess, normal.Normal)
def _kl_gp_normal(gp, n, name=None):
  """Calculate the batched KL divergence KL(gp || n).

  Args:
    gp: instance of a GaussianProcess distribution object.
    n: instance of a Normal distribution object.
    name: (optional) Name to use for created operations.
      default is 'kl_gp_normal'.

  Returns:
    Batchwise KL(gp || n)
  """
  return _kl_gp_compatible(gp, n, name or 'kl_gp_normal')

@kullback_leibler.RegisterKL(
    GaussianProcess, mvn_linear_operator.MultivariateNormalLinearOperator)
def _kl_gp_mvn(gp, mvn, name=None):
  """Calculate the batched KL divergence KL(gp || mvn).

  Args:
    gp: instance of a GaussianProcess distribution object.
    mvn: instance of a multivariate Normal distribution object (any subclass of
      MultivariateNormalLinearOperator)
    name: (optional) Name to use for created operations.
      default is 'kl_gp_mvn'.

  Returns:
    Batchwise KL(gp || mvn)
  """
  return _kl_gp_compatible(gp, mvn, name or 'kl_gp_mvn')

@kullback_leibler.RegisterKL(normal.Normal, GaussianProcess)
def _kl_normal_gp(n, gp, name=None):
  """Calculate the batched KL divergence KL(gp || n).

  Args:
    n: instance of a Normal distribution object.
    gp: instance of a GaussianProcess distribution object.
    name: (optional) Name to use for created operations.
      default is 'kl_normal_gp'.

  Returns:
    Batchwise KL(n || gp)
  """
  return _kl_compatible_gp(n, gp, name or 'kl_normal_gp')

@kullback_leibler.RegisterKL(
    mvn_linear_operator.MultivariateNormalLinearOperator, GaussianProcess)
def _kl_mvn_gp(mvn, gp, name=None):
  """Calculate the batched KL divergence KL(mvn || gp).

  Args:
    mvn: instance of a multivariate Normal distribution object (any subclass of
      MultivariateNormalLinearOperator)
    gp: instance of a GaussianProcess distribution object.
    name: (optional) Name to use for created operations.
      default is 'kl_mvn_gp'.

  Returns:
    Batchwise KL(mvn || gp)
  """
  return _kl_compatible_gp(mvn, gp, name or 'kl_mvn_gp')


# 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.
# ============================================================================
"""Quantized distribution."""

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 as distributions
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

__all__ = ["QuantizedDistribution"]

def _logsum_expbig_minus_expsmall(big, small):
  """Stable evaluation of `Log[exp{big} - exp{small}]`.

  To work correctly, we should have the pointwise relation:  `small <= big`.

  Args:
    big: Floating-point `Tensor`
    small: Floating-point `Tensor` with same `dtype` as `big` and broadcastable
      shape.

  Returns:
    `Tensor` of same `dtype` of `big` and broadcast shape.
  """
  with tf.name_scope("logsum_expbig_minus_expsmall"):
    return tf.math.log1p(-tf.exp(small - big)) + big

_prob_base_note = """
For whole numbers `y`,

```
P[Y = y] := P[X <= low],  if y == low,
         := P[X > high - 1],  y == high,
         := 0, if j < low or y > high,
         := P[y - 1 < X <= y],  all other y.
```

"""

_prob_note = _prob_base_note + """
The base distribution's `cdf` method must be defined on `y - 1`. If the
base distribution has a `survival_function` method, results will be more
accurate for large values of `y`, and in this case the `survival_function` must
also be defined on `y - 1`.
"""

_log_prob_note = _prob_base_note + """
The base distribution's `log_cdf` method must be defined on `y - 1`. If the
base distribution has a `log_survival_function` method results will be more
accurate for large values of `y`, and in this case the `log_survival_function`
must also be defined on `y - 1`.
"""

_cdf_base_note = """

For whole numbers `y`,

```
cdf(y) := P[Y <= y]
        = 1, if y >= high,
        = 0, if y < low,
        = P[X <= y], otherwise.
```

Since `Y` only has mass at whole numbers, `P[Y <= y] = P[Y <= floor(y)]`.
This dictates that fractional `y` are first floored to a whole number, and
then above definition applies.
"""

_cdf_note = _cdf_base_note + """
The base distribution's `cdf` method must be defined on `y - 1`.
"""

_log_cdf_note = _cdf_base_note + """
The base distribution's `log_cdf` method must be defined on `y - 1`.
"""

_sf_base_note = """

For whole numbers `y`,

```
survival_function(y) := P[Y > y]
                      = 0, if y >= high,
                      = 1, if y < low,
                      = P[X <= y], otherwise.
```

Since `Y` only has mass at whole numbers, `P[Y <= y] = P[Y <= floor(y)]`.
This dictates that fractional `y` are first floored to a whole number, and
then above definition applies.
"""

_sf_note = _sf_base_note + """
The base distribution's `cdf` method must be defined on `y - 1`.
"""

_log_sf_note = _sf_base_note + """
The base distribution's `log_cdf` method must be defined on `y - 1`.
"""

class QuantizedDistribution(distributions.Distribution):
  """Distribution representing the quantization `Y = ceiling(X)`.

  #### Definition in Terms of Sampling

  ```
  1. Draw X
  2. Set Y <-- ceiling(X)
  3. If Y < low, reset Y <-- low
  4. If Y > high, reset Y <-- high
  5. Return Y
  ```

  #### Definition in Terms of the Probability Mass Function

  Given scalar random variable `X`, we define a discrete random variable `Y`
  supported on the integers as follows:

  ```
  P[Y = j] := P[X <= low],  if j == low,
           := P[X > high - 1],  j == high,
           := 0, if j < low or j > high,
           := P[j - 1 < X <= j],  all other j.
  ```

  Conceptually, without cutoffs, the quantization process partitions the real
  line `R` into half open intervals, and identifies an integer `j` with the
  right endpoints:

  ```
  R = ... (-2, -1](-1, 0](0, 1](1, 2](2, 3](3, 4] ...
  j = ...      -1      0     1     2     3     4  ...
  ```

  `P[Y = j]` is the mass of `X` within the `jth` interval.
  If `low = 0`, and `high = 2`, then the intervals are redrawn
  and `j` is re-assigned:

  ```
  R = (-infty, 0](0, 1](1, infty)
  j =          0     1     2
  ```

  `P[Y = j]` is still the mass of `X` within the `jth` interval.

  #### Examples

  We illustrate a mixture of discretized logistic distributions
  [(Salimans et al., 2017)][1]. This is used, for example, for capturing 16-bit
  audio in WaveNet [(van den Oord et al., 2017)][2]. The values range in
  a 1-D integer domain of `[0, 2**16-1]`, and the discretization captures
  `P(x - 0.5 < X <= x + 0.5)` for all `x` in the domain excluding the endpoints.
  The lowest value has probability `P(X <= 0.5)` and the highest value has
  probability `P(2**16 - 1.5 < X)`.

  Below we assume a `wavenet` function. It takes as `input` right-shifted audio
  samples of shape `[..., sequence_length]`. It returns a real-valued tensor of
  shape `[..., num_mixtures * 3]`, i.e., each mixture component has a `loc` and
  `scale` parameter belonging to the logistic distribution, and a `logits`
  parameter determining the unnormalized probability of that component.

  ```python
  tfd = tfp.distributions
  tfb = tfp.bijectors

  net = wavenet(inputs)
  loc, unconstrained_scale, logits = tf.split(net,
                                              num_or_size_splits=3,
                                              axis=-1)
  scale = tf.nn.softplus(unconstrained_scale)

  # Form mixture of discretized logistic distributions. Note we shift the
  # logistic distribution by -0.5. This lets the quantization capture "rounding"
  # intervals, `(x-0.5, x+0.5]`, and not "ceiling" intervals, `(x-1, x]`.
  discretized_logistic_dist = tfd.QuantizedDistribution(
      distribution=tfd.TransformedDistribution(
          distribution=tfd.Logistic(loc=loc, scale=scale),
          bijector=tfb.AffineScalar(shift=-0.5)),
      low=0.,
      high=2**16 - 1.)
  mixture_dist = tfd.MixtureSameFamily(
      mixture_distribution=tfd.Categorical(logits=logits),
      components_distribution=discretized_logistic_dist)

  neg_log_likelihood = -tf.reduce_sum(mixture_dist.log_prob(targets))
  train_op = tf.train.AdamOptimizer().minimize(neg_log_likelihood)
  ```

  After instantiating `mixture_dist`, we illustrate maximum likelihood by
  calculating its log-probability of audio samples as `target` and optimizing.

  #### References

  [1]: Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma.
       PixelCNN++: Improving the PixelCNN with discretized logistic mixture
       likelihood and other modifications.
       _International Conference on Learning Representations_, 2017.
       https://arxiv.org/abs/1701.05517
  [2]: Aaron van den Oord et al. Parallel WaveNet: Fast High-Fidelity Speech
       Synthesis. _arXiv preprint arXiv:1711.10433_, 2017.
       https://arxiv.org/abs/1711.10433
  """

  def __init__(self,
               distribution,
               low=None,
               high=None,
               validate_args=False,
               name="QuantizedDistribution"):
    """Construct a Quantized Distribution representing `Y = ceiling(X)`.

    Some properties are inherited from the distribution defining `X`. Example:
    `allow_nan_stats` is determined for this `QuantizedDistribution` by reading
    the `distribution`.

    Args:
      distribution:  The base distribution class to transform. Typically an
        instance of `Distribution`.
      low: `Tensor` with same `dtype` as this distribution and shape
        able to be added to samples. Should be a whole number. Default `None`.
        If provided, base distribution's `prob` should be defined at
        `low`.
      high: `Tensor` with same `dtype` as this distribution and shape
        able to be added to samples. Should be a whole number. Default `None`.
        If provided, base distribution's `prob` should be defined at
        `high - 1`.
        `high` must be strictly greater than `low`.
      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.
      name: Python `str` name prefixed to Ops created by this class.

    Raises:
      TypeError: If `dist_cls` is not a subclass of
          `Distribution` or continuous.
      NotImplementedError:  If the base distribution does not implement `cdf`.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      self._dist = distribution

      if low is not None:
        low = tf.convert_to_tensor(
            value=low, name="low", dtype=distribution.dtype)
      if high is not None:
        high = tf.convert_to_tensor(
            value=high, name="high", dtype=distribution.dtype)
      dtype_util.assert_same_float_dtype(
          tensors=[self.distribution, low, high])

      # We let QuantizedDistribution access _graph_parents since this class is
      # more like a baseclass.
      graph_parents = self._dist._graph_parents  # pylint: disable=protected-access

      checks = []
      if validate_args and low is not None and high is not None:
        message = "low must be strictly less than high."
        checks.append(assert_util.assert_less(low, high, message=message))
      self._validate_args = validate_args  # self._check_integer uses this.
      with tf.control_dependencies(checks if validate_args else []):
        if low is not None:
          self._low = self._check_integer(low)
          graph_parents += [self._low]
        else:
          self._low = None
        if high is not None:
          self._high = self._check_integer(high)
          graph_parents += [self._high]
        else:
          self._high = None

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

  @property
  def distribution(self):
    """Base distribution, p(x)."""
    return self._dist

  @property
  def low(self):
    """Lowest value that quantization returns."""
    return self._low

  @property
  def high(self):
    """Highest value that quantization returns."""
    return self._high

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

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

  def _event_shape_tensor(self):
    return self.distribution.event_shape_tensor()

  def _event_shape(self):
    return self.distribution.event_shape

  def _sample_n(self, n, seed=None):
    low = self._low
    high = self._high
    with tf.name_scope("transform"):
      n = tf.convert_to_tensor(value=n, name="n")
      x_samps = self.distribution.sample(n, seed=seed)
      ones = tf.ones_like(x_samps)

      # Snap values to the intervals (j - 1, j].
      result_so_far = tf.math.ceil(x_samps)

      if low is not None:
        result_so_far = tf.where(result_so_far < low, low * ones, result_so_far)

      if high is not None:
        result_so_far = tf.where(result_so_far > high, high * ones,
                                 result_so_far)

      return result_so_far

  @distribution_util.AppendDocstring(_log_prob_note)
  def _log_prob(self, y):
    if not hasattr(self.distribution, "_log_cdf"):
      raise NotImplementedError(
          "'log_prob' not implemented unless the base distribution implements "
          "'log_cdf'")
    y = self._check_integer(y)
    try:
      return self._log_prob_with_logsf_and_logcdf(y)
    except NotImplementedError:
      return self._log_prob_with_logcdf(y)

  def _log_prob_with_logcdf(self, y):
    return _logsum_expbig_minus_expsmall(self.log_cdf(y), self.log_cdf(y - 1))

  def _log_prob_with_logsf_and_logcdf(self, y):
    """Compute log_prob(y) using log survival_function and cdf together."""
    # There are two options that would be equal if we had infinite precision:
    # Log[ sf(y - 1) - sf(y) ]
    #   = Log[ exp{logsf(y - 1)} - exp{logsf(y)} ]
    # Log[ cdf(y) - cdf(y - 1) ]
    #   = Log[ exp{logcdf(y)} - exp{logcdf(y - 1)} ]
    logsf_y = self.log_survival_function(y)
    logsf_y_minus_1 = self.log_survival_function(y - 1)
    logcdf_y = self.log_cdf(y)
    logcdf_y_minus_1 = self.log_cdf(y - 1)

    # Important:  Here we use select in a way such that no input is inf, this
    # prevents the troublesome case where the output of select can be finite,
    # but the output of grad(select) will be NaN.

    # In either case, we are doing Log[ exp{big} - exp{small} ]
    # We want to use the sf items precisely when we are on the right side of the
    # median, which occurs when logsf_y < logcdf_y.
    big = tf.where(logsf_y < logcdf_y, logsf_y_minus_1, logcdf_y)
    small = tf.where(logsf_y < logcdf_y, logsf_y, logcdf_y_minus_1)

    return _logsum_expbig_minus_expsmall(big, small)

  @distribution_util.AppendDocstring(_prob_note)
  def _prob(self, y):
    if not hasattr(self.distribution, "_cdf"):
      raise NotImplementedError(
          "'prob' not implemented unless the base distribution implements "
          "'cdf'")
    y = self._check_integer(y)
    try:
      return self._prob_with_sf_and_cdf(y)
    except NotImplementedError:
      return self._prob_with_cdf(y)

  def _prob_with_cdf(self, y):
    return self.cdf(y) - self.cdf(y - 1)

  def _prob_with_sf_and_cdf(self, y):
    # There are two options that would be equal if we had infinite precision:
    # sf(y - 1) - sf(y)
    # cdf(y) - cdf(y - 1)
    sf_y = self.survival_function(y)
    sf_y_minus_1 = self.survival_function(y - 1)
    cdf_y = self.cdf(y)
    cdf_y_minus_1 = self.cdf(y - 1)

    # sf_prob has greater precision iff we're on the right side of the median.
    return tf.where(
        sf_y < cdf_y,  # True iff we're on the right side of the median.
        sf_y_minus_1 - sf_y,
        cdf_y - cdf_y_minus_1)

  @distribution_util.AppendDocstring(_log_cdf_note)
  def _log_cdf(self, y):
    low = self._low
    high = self._high

    # Recall the promise:
    # cdf(y) := P[Y <= y]
    #         = 1, if y >= high,
    #         = 0, if y < low,
    #         = P[X <= y], otherwise.

    # P[Y <= j] = P[floor(Y) <= j] since mass is only at integers, not in
    # between.
    j = tf.floor(y)

    result_so_far = self.distribution.log_cdf(j)

    # Broadcast, because it's possible that this is a single distribution being
    # evaluated on a number of samples, or something like that.
    j += tf.zeros_like(result_so_far)

    # Re-define values at the cutoffs.
    if low is not None:
      neg_inf = -np.inf * tf.ones_like(result_so_far)
      result_so_far = tf.where(j < low, neg_inf, result_so_far)
    if high is not None:
      result_so_far = tf.where(j >= high, tf.zeros_like(result_so_far),
                               result_so_far)

    return result_so_far

  @distribution_util.AppendDocstring(_cdf_note)
  def _cdf(self, y):
    low = self._low
    high = self._high

    # Recall the promise:
    # cdf(y) := P[Y <= y]
    #         = 1, if y >= high,
    #         = 0, if y < low,
    #         = P[X <= y], otherwise.

    # P[Y <= j] = P[floor(Y) <= j] since mass is only at integers, not in
    # between.
    j = tf.floor(y)

    # P[X <= j], used when low < X < high.
    result_so_far = self.distribution.cdf(j)

    # Broadcast, because it's possible that this is a single distribution being
    # evaluated on a number of samples, or something like that.
    j += tf.zeros_like(result_so_far)

    # Re-define values at the cutoffs.
    if low is not None:
      result_so_far = tf.where(j < low, tf.zeros_like(result_so_far),
                               result_so_far)
    if high is not None:
      result_so_far = tf.where(j >= high, tf.ones_like(result_so_far),
                               result_so_far)

    return result_so_far

  @distribution_util.AppendDocstring(_log_sf_note)
  def _log_survival_function(self, y):
    low = self._low
    high = self._high

    # Recall the promise:
    # survival_function(y) := P[Y > y]
    #                       = 0, if y >= high,
    #                       = 1, if y < low,
    #                       = P[X > y], otherwise.

    # P[Y > j] = P[ceiling(Y) > j] since mass is only at integers, not in
    # between.
    j = tf.math.ceil(y)

    # P[X > j], used when low < X < high.
    result_so_far = self.distribution.log_survival_function(j)

    # Broadcast, because it's possible that this is a single distribution being
    # evaluated on a number of samples, or something like that.
    j += tf.zeros_like(result_so_far)

    # Re-define values at the cutoffs.
    if low is not None:
      result_so_far = tf.where(j < low, tf.zeros_like(result_so_far),
                               result_so_far)
    if high is not None:
      neg_inf = -np.inf * tf.ones_like(result_so_far)
      result_so_far = tf.where(j >= high, neg_inf, result_so_far)

    return result_so_far

  @distribution_util.AppendDocstring(_sf_note)
  def _survival_function(self, y):
    low = self._low
    high = self._high

    # Recall the promise:
    # survival_function(y) := P[Y > y]
    #                       = 0, if y >= high,
    #                       = 1, if y < low,
    #                       = P[X > y], otherwise.

    # P[Y > j] = P[ceiling(Y) > j] since mass is only at integers, not in
    # between.
    j = tf.math.ceil(y)

    # P[X > j], used when low < X < high.
    result_so_far = self.distribution.survival_function(j)

    # Broadcast, because it's possible that this is a single distribution being
    # evaluated on a number of samples, or something like that.
    j += tf.zeros_like(result_so_far)

    # Re-define values at the cutoffs.
    if low is not None:
      result_so_far = tf.where(j < low, tf.ones_like(result_so_far),
                               result_so_far)
    if high is not None:
      result_so_far = tf.where(j >= high, tf.zeros_like(result_so_far),
                               result_so_far)

    return result_so_far

  def _check_integer(self, value):
    with tf.name_scope("check_integer"):
      value = tf.convert_to_tensor(value=value, name="value")
      if not self.validate_args:
        return value
      dependencies = [distribution_util.assert_integer_form(
          value, message="value has non-integer components.")]
      return distribution_util.with_dependencies(dependencies, value)


# 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.
# ============================================================================
"""Implementation of the NeuTra Kernel."""

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

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

from tensorflow_probability.python.internal import dtype_util

tfb = tfp.bijectors
tfd = tfp.distributions

__all__ = ['NeuTra', 'make_iaf_stack']

def make_iaf_stack(total_event_size,
                   num_hidden_layers=2,
                   seed=None,
                   dtype=tf.float32):
  """Creates an stacked IAF bijector.

  This bijector operates on vector-valued events.

  Args:
    total_event_size: Number of dimensions to operate over.
    num_hidden_layers: How many hidden layers to use in each IAF.
    seed: Random seed for the initializers.
    dtype: DType for the variables.

  Returns:
    bijector: The created bijector.
  """

  seed = tfd.SeedStream(seed, 'make_iaf_stack')

  def make_iaf():
    """Create an IAF."""
    initializer = tf.compat.v2.keras.initializers.VarianceScaling(
        2 * 0.01, seed=seed() % (2**31 - 1))

    made = tfb.AutoregressiveLayer(
        params=2,
        event_shape=[total_event_size],
        hidden_units=[total_event_size] * num_hidden_layers,
        activation=tf.nn.elu,
        kernel_initializer=initializer,
        dtype=dtype)

    def shift_and_scale(x):
      # TODO(siege): Something is losing the static shape.
      x.set_shape(
          x.shape.merge_with([None] * (x.shape.ndims - 1) + [total_event_size]))
      return tf.unstack(made(x), num=2, axis=-1)

    return tfb.Invert(tfb.MaskedAutoregressiveFlow(shift_and_scale))

  def make_swap():
    """Create an swap."""
    permutation = list(reversed(range(total_event_size)))
    return tfb.Permute(permutation)

  bijector = make_iaf()
  bijector = make_swap()(bijector)
  bijector = make_iaf()(bijector)
  bijector = make_swap()(bijector)
  bijector = make_iaf()(bijector)
  bijector = make_swap()(bijector)

  return bijector

class NeuTra(tfp.mcmc.TransitionKernel):
  """The experimental NeuTra kernel.

  Warning: This kernel is experimental. Default arguments and their
  interpreations will very likely change.

  This kernel does not work in TF1 graph mode.

  This implements a transition kernel that implements the NeuTra MCMC
  [(Hoffman et al., 2019)][1]. It operates by learning a neural pre-conditioner
  (bijector) inside the `bootstrap_results` method, and then using it inside
  `one_step` to accelerate HMC. The same bijector is used to initialize the
  state: i.e. the `current_state` as passed to `tfp.mcmc.sample_chain` is
  ignored except to extract the number of chains to sample in parallel.

  This kernel performs step-size adaptation and picks an automatic trajectory
  length for its internal HMC kernel.

  If your problem has constrained support, specify this support via the
  `unconstraining_bijector` argument. This argument is interpreted such that the
  image of the forward transformation of that bijector matches the support of
  your problem. E.g. if one of your random variables is positive, you could use
  the `tfb.Softplus` bijector.

  Since this is still experimental, we provide some facilities to debug the
  bijector training via the `train_debug_fn`, where you can monitor the training
  progress. In practice, it may be prudent to run the training process multiple
  times to verify that the variational approximation is stable. If it is not,
  you can attempt to increase the expressiveness of the bijector via the
  `trainable_bijector_fn` argument. Currently the bijector operates by
  flattening the entire problem state into one vector. This is seamless to the
  user, but is important to remember when designing a bijector. The flattening
  will likely be removed in the future.

  `_flattened_variational_distribution` returns the (flattened) distribution
  resulting from a standard normal being pushed through the bijector. This is
  useful for debugging, as this can be compared to the MCMC chain and your prior
  preconception of what the target distribution should look like.

  Additionally, you can examine the `inner_results.transformed_state` inside the
  kernel results. By construction, for a well-fitted bijector this should
  resemble a standard normal. This kernel assumes that you have a well-fitted
  bijector, so if samples in the transformed space do not look normal, then this
  kernel is not operating efficiently.

  ### Examples

  Sampling from a multivariate log-normal distribution.

  ```python

  target_dist = tfd.MultivariateNormalTriL(scale_tril=[[1., 0.], [2, 1.]])
  target_dist = tfb.Exp()(target_dist)

  num_chains = 64
  state_shape = 2
  num_burnin_steps = 1000

  kernel = NeuTra(
      target_log_prob_fn=target_dist.log_prob,
      state_shape=state_shape,
      num_step_size_adaptation_steps=int(0.8 * num_burnin_steps),
      unconstraining_bijector=tfb.Exp())

  chain = tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=1000,
      current_state=tf.ones([num_chains, 2]),
      kernel=kernel,
      trace_fn=None)
  ```

  #### References

  [1]: Hoffman, M., Sountsov, P., Dillon, J. V., Langmore, I., Tran, D., &
       Vasudevan, S. (2019). NeuTra-lizing Bad Geometry in Hamiltonian Monte
       Carlo Using Neural Transport. http://arxiv.org/abs/1903.03704
  """

  def __init__(self,
               target_log_prob_fn,
               state_shape,
               num_step_size_adaptation_steps,
               unconstraining_bijector=None,
               trainable_bijector_fn=make_iaf_stack,
               learning_rate=1e-2,
               train_batch_size=4096,
               num_train_steps=5000,
               train_debug_fn=None,
               seed=None):
    """Creates the 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.
      state_shape: A Python list, or list of Python lists (or `TensorShape`s)
        that describes the shape of the problem state. Must be the same
        structure as `current_state`. All shapes must be fully defined.
      num_step_size_adaptation_steps: Number of steps to use for step size
        adaptation. See `tfp.mcmc.SimpleStepSizeAdaptation`.
      unconstraining_bijector: A bijector or list of bijectors that go from
        unconstrained space to the support of the corresponding random variable.
        Must be the same structure as `current_state`.
      trainable_bijector_fn: Creates a trainable, vector-valued event bijector.
        Must be a callable with signature: `(total_event_size, seed) ->
          bijector`, where `total_event_size` is the size of the event.
      learning_rate: Base learning rate to use for training the bijector.
        Internally, learning rate decay is used to stabilize learning.
      train_batch_size: Batch size to use for training the bijector.
      num_train_steps: Number of training steps to train the bijector.
      train_debug_fn: A callable with signature `(NeuTra, step, loss)` called
        for every training step. The first argument is this instance, and `step`
        is the current training step.
      seed: A seed for reproducibility.
    """
    self._parameters = dict(
        target_log_prob_fn=target_log_prob_fn,
        state_shape=state_shape,
        num_step_size_adaptation_steps=num_step_size_adaptation_steps,
        unconstraining_bijector=unconstraining_bijector,
        trainable_bijector_fn=trainable_bijector_fn,
        learning_rate=learning_rate,
        train_batch_size=train_batch_size,
        num_train_steps=num_train_steps,
        train_debug_fn=train_debug_fn,
        seed=seed)

    self._state_shape = tf.nest.map_structure(tf.TensorShape, state_shape)

    if unconstraining_bijector is None:
      unconstraining_bijector = tf.nest.map_structure(lambda _: tfb.Identity(),
                                                      self.state_shape)
    self._unconstraining_bijector = unconstraining_bijector

    def _make_reshaped_bijector(b, s):
      return tfb.Reshape(
          event_shape_in=s, event_shape_out=[s.num_elements()])(b)(
              tfb.Reshape(event_shape_out=b.inverse_event_shape(s)))

    # This converts the `unconstraining_bijector` to work on flattened state
    # parts.
    reshaped_bijector = tf.nest.map_structure(_make_reshaped_bijector,
                                              unconstraining_bijector,
                                              self.state_shape)

    blockwise_bijector = tfb.Blockwise(
        bijectors=tf.nest.flatten(reshaped_bijector),
        block_sizes=tf.nest.flatten(
            tf.nest.map_structure(lambda s: s.num_elements(),
                                  self.state_shape)))

    dtype = dtype_util.common_dtype([learning_rate], preferred_dtype=tf.float32)
    self._dtype = dtype

    trainable_bijector = trainable_bijector_fn(
        self._total_event_size, seed=seed, dtype=dtype)
    self._trainable_bijector = trainable_bijector

    self._bijector = blockwise_bijector(trainable_bijector)

    def flattened_target_log_prob(flat_state):
      state = self._unflatten_state(flat_state)
      if isinstance(state, (list, tuple)):
        return target_log_prob_fn(*state)
      else:
        return target_log_prob_fn(state)

    self._flattened_target_log_prob_val = flattened_target_log_prob

    kernel = tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=flattened_target_log_prob,
        step_size=1.,
        num_leapfrog_steps=self._num_leapfrog_steps(1.),
        seed=seed if not tf.executing_eagerly() else None)
    kernel = tfp.mcmc.TransformedTransitionKernel(kernel, self._bijector)
    kernel = tfp.mcmc.SimpleStepSizeAdaptation(
        kernel,
        num_adaptation_steps=num_step_size_adaptation_steps,
        target_accept_prob=np.array(0.9, dtype.as_numpy_dtype()))
    self._kernel = kernel

  @property
  def trainable_bijector(self):
    return self._trainable_bijector

  @property
  def unconstraining_bijector(self):
    return self._unconstraining_bijector

  @property
  def state_shape(self):
    return self._state_shape

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

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

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

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

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

  @property
  def _total_event_size(self):
    return sum(
        tf.nest.flatten(
            tf.nest.map_structure(
                lambda b, s: b.inverse_event_shape(s).num_elements(),
                self.unconstraining_bijector, self.state_shape)))

  def _num_leapfrog_steps(self, step_size):
    step_size = tf.convert_to_tensor(value=step_size)
    trajectory_length = np.float32(self._total_event_size)**0.25
    return tf.cast(tf.math.ceil(trajectory_length / step_size), dtype=tf.int32)

  def _flatten_state(self, state):
    state_parts = tf.nest.flatten(state)
    flat_state_shapes = tf.nest.flatten(self.state_shape)
    batch_shape = tf.shape(input=state_parts[0])[:-flat_state_shapes[0].ndims]
    flat_shape = tf.concat([batch_shape, [-1]], -1)
    flat_state_parts = tf.nest.map_structure(
        lambda s: tf.reshape(s, flat_shape), state_parts)
    return tf.concat(flat_state_parts, -1)

  def _unflatten_state(self, flat_state):
    state_parts = tf.split(
        flat_state,
        [s.num_elements() for s in tf.nest.flatten(self.state_shape)],
        axis=-1)
    batch_shape = tf.shape(input=flat_state)[:-1]

    state = tf.nest.pack_sequence_as(self.state_shape, state_parts)
    return tf.nest.map_structure(
        lambda part, s: tf.reshape(part, tf.concat([batch_shape, s], 0)), state,
        self.state_shape)

  def is_calibrated(self):
    return True

  def _flattened_variational_distribution(self):
    base = tfd.MultivariateNormalDiag(
        loc=tf.zeros(self._total_event_size, dtype=self._dtype))
    return self._bijector(base)

  @property
  def _flattened_target_log_prob(self):
    return self._flattened_target_log_prob_val

  @tf.function(autograph=False)
  def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of NeuTra.

    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.
    """

    @tfp.mcmc.internal.util.make_innermost_setter
    def set_num_leapfrog_steps(kernel_results, num_leapfrog_steps):
      return kernel_results._replace(
          accepted_results=kernel_results.accepted_results._replace(
              num_leapfrog_steps=num_leapfrog_steps))

    step_size = previous_kernel_results.new_step_size
    previous_kernel_results = set_num_leapfrog_steps(
        previous_kernel_results, self._num_leapfrog_steps(step_size))

    new_state, kernel_results = self._kernel.one_step(
        self._flatten_state(current_state), previous_kernel_results)
    return self._unflatten_state(new_state), kernel_results

  def bootstrap_results(self, state):
    """Trains the bijector and creates initial `previous_kernel_results`.

    The supplied `state` is only used to determine the number of chains to run
    in parallel_iterations

    Args:
      state: `Tensor` or Python `list` of `Tensor`s representing the initial
        state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*state))`.

    Returns:
      kernel_results: Instance of
        `UncalibratedHamiltonianMonteCarloKernelResults` inside
        `MetropolisHastingsResults` inside `TransformedTransitionKernelResults`
        inside `SimpleStepSizeAdaptationResults`.
    """

    def loss():
      q = self._flattened_variational_distribution()
      # TODO(siege): How to seed this?
      samples = q.sample(self.train_batch_size)
      return tf.reduce_mean(
          input_tensor=q.log_prob(samples) -
          self._flattened_target_log_prob(samples),
          axis=-1)

    lr = tf.convert_to_tensor(value=self.learning_rate, dtype=self._dtype)
    dtype = lr.dtype

    learning_rate = tf.compat.v2.optimizers.schedules.PiecewiseConstantDecay(
        list(self.num_train_steps *
             np.array([0.2, 0.8]).astype(dtype.as_numpy_dtype())),
        [lr, lr * 0.1, lr * 0.01])

    opt = tf.compat.v2.optimizers.Adam(learning_rate)

    @tf.function(autograph=False)
    def train_step():
      with tf.GradientTape() as tape:
        loss_val = loss()
      vals = tape.watched_variables()
      grads = tape.gradient(loss_val, vals)
      grads_and_vals = list(zip(grads, vals))
      opt.apply_gradients(grads_and_vals)
      return loss_val

    for step in range(self.num_train_steps):
      loss_val = train_step()
      tf.debugging.assert_all_finite(
          loss_val, 'NeuTra loss is NaN at step {}'.format(step))
      if self.train_debug_fn:
        # pylint: disable=not-callable
        self.train_debug_fn(self, step, loss_val)

    state_parts = tf.nest.flatten(state)
    flat_state_shapes = tf.nest.flatten(self.state_shape)
    batch_shape = tf.shape(input=state_parts[0])[:-flat_state_shapes[0].ndims]

    return self._kernel.bootstrap_results(
        self._flattened_variational_distribution().sample(
            batch_shape, seed=self.seed))