# 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.
# ============================================================================
"""Bijector base."""

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

import abc
import collections
import contextlib
import re
import weakref

# Dependency imports
import numpy as np
import six
import tensorflow as tf

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 tensorshape_util
from tensorflow.python.util import deprecation  # pylint: disable=g-direct-tensorflow-import

__all__ = [
    "Bijector",
    "ConditionalBijector",
]

def _get_current_graph():
  if tf.executing_eagerly():
    return None
  return tf.compat.v1.get_default_graph()

class _Mapping(
    collections.namedtuple("_Mapping", ["x", "y", "ildj", "kwargs"])):
  """Helper class to make it easier to manage caching in `Bijector`."""

  def __new__(cls, x=None, y=None, ildj=None, kwargs=None):
    """Custom __new__ so namedtuple items have defaults.

    Args:
      x: `Tensor` or None. Input to forward; output of inverse.
      y: `Tensor` or None. Input to inverse; output of forward.
      ildj: `Tensor`. This is the (un-reduce_sum'ed) inverse log det jacobian.
      kwargs: Python dictionary. Extra args supplied to forward/inverse/etc
        functions.

    Returns:
      mapping: New instance of _Mapping.
    """
    return super(_Mapping, cls).__new__(cls, x, y, ildj, kwargs)

  @property
  def subkey(self):
    """Returns subkey used for caching (nested under either `x` or `y`)."""
    return self._deep_tuple(self.kwargs)

  def merge(self, x=None, y=None, ildj=None, kwargs=None, mapping=None):
    """Returns new _Mapping with args merged with self.

    Args:
      x: `Tensor` or None. Input to forward; output of inverse.
      y: `Tensor` or None. Input to inverse; output of forward.
      ildj: `Tensor`. This is the (un-reduce_sum'ed) inverse log det jacobian.
      kwargs: Python dictionary. Extra args supplied to forward/inverse/etc
        functions.
      mapping: Instance of _Mapping to merge. Can only be specified if no other
        arg is specified.

    Returns:
      mapping: New instance of `_Mapping` which has inputs merged with self.

    Raises:
      ValueError: if mapping and any other arg is not `None`.
    """
    if mapping is None:
      mapping = _Mapping(x=x, y=y, ildj=ildj, kwargs=kwargs)
    elif any(arg is not None for arg in [x, y, ildj, kwargs]):
      raise ValueError("Cannot simultaneously specify mapping and individual "
                       "arguments.")

    return _Mapping(
        x=self._merge(self.x, mapping.x),
        y=self._merge(self.y, mapping.y),
        ildj=self._merge(self.ildj, mapping.ildj),
        kwargs=self._merge(self.kwargs, mapping.kwargs, use_equals=True))

  def remove(self, field):
    """To support weak referencing, removes cache key from the cache value."""
    return _Mapping(
        x=None if field == "x" else self.x,
        y=None if field == "y" else self.y,
        ildj=self.ildj,
        kwargs=self.kwargs)

  def _merge(self, old, new, use_equals=False):
    """Helper to merge which handles merging one value."""
    if old is None:
      return new
    if new is None:
      return old
    if (old == new) if use_equals else (old is new):
      return old
    raise ValueError("Incompatible values: %s != %s" % (old, new))

  def _deep_tuple(self, x):
    """Converts nested `tuple`, `list`, or `dict` to nested `tuple`."""
    if isinstance(x, dict):
      return self._deep_tuple(tuple(sorted(x.items())))
    elif isinstance(x, (list, tuple)):
      return tuple(map(self._deep_tuple, x))

    return x

class WeakKeyDefaultDict(dict):
  """`WeakKeyDictionary` which always adds `defaultdict(dict)` in getitem."""

  # Q:Why not subclass `collections.defaultdict`?
  # Subclassing collections.defaultdict means we have a more complicated `repr`,
  # `str` which makes debugging the bijector cache more tedious. Additionally it
  # means we need to think about passing through __init__ args but manually
  # specifying the `default_factory`. That is, just overriding `__missing__`
  # ends up being a lot cleaner.

  # Q:Why not subclass `weakref.WeakKeyDictionary`?
  # `weakref.WeakKeyDictionary` has an even worse `repr`, `str` than
  # collections.defaultdict. Plus, since we want explicit control over how the
  # keys are created we need to override __getitem__ which is the only feature
  # of `weakref.WeakKeyDictionary` we're using.

  # This is the "WeakKey" part.
  def __getitem__(self, key):
    weak_key = HashableWeakRef(key, lambda w: self.pop(w, None))
    return super(WeakKeyDefaultDict, self).__getitem__(weak_key)

  # This is the "DefaultDict" part.
  def __missing__(self, key):
    assert isinstance(key, HashableWeakRef)  # Can't happen.
    return super(WeakKeyDefaultDict, self).setdefault(key, {})

  # Everything that follows is only useful to help make debugging easier.

  def __contains__(self, key):
    return super(WeakKeyDefaultDict, self).__contains__(HashableWeakRef(key))

  # We don't want mutation except through __getitem__.

  def __setitem__(self, *args, **kwargs):
    raise NotImplementedError()

  def update(self, *args, **kwargs):
    raise NotImplementedError()

  def setdefault(self, *args, **kwargs):
    raise NotImplementedError()

class HashableWeakRef(weakref.ref):
  """weakref.ref which makes np.array objects hashable."""

  def __hash__(self):
    x = self()
    if not isinstance(x, np.ndarray):
      return hash(x)
    # Note: The following logic can never be reached by the public API because
    # the bijector base class always calls `convert_to_tensor` before accessing
    # the cache.
    x.flags.writeable = False
    return hash(str(x.__array_interface__) + str(id(x)))

  def __repr__(self):
    return repr(self())

  def __str__(self):
    return str(self())

  def __eq__(self, other):
    x = self()
    if isinstance(x, np.ndarray):
      y = other()
      return (isinstance(y, np.ndarray) and
              x.__array_interface__ == y.__array_interface__ and
              id(x) == id(y))
    return super(HashableWeakRef, self).__eq__(other)

@six.add_metaclass(abc.ABCMeta)
class Bijector(object):
  r"""Interface for transformations of a `Distribution` sample.

  Bijectors can be used to represent any differentiable and injective
  (one to one) function defined on an open subset of `R^n`.  Some non-injective
  transformations are also supported (see "Non Injective Transforms" below).

  #### Mathematical Details

  A `Bijector` implements a [smooth covering map](
  https://en.wikipedia.org/wiki/Local_diffeomorphism), i.e., a local
  diffeomorphism such that every point in the target has a neighborhood evenly
  covered by a map ([see also](
  https://en.wikipedia.org/wiki/Covering_space#Covering_of_a_manifold)).
  A `Bijector` is used by `TransformedDistribution` but can be generally used
  for transforming a `Distribution` generated `Tensor`. A `Bijector` is
  characterized by three operations:

  1. Forward

     Useful for turning one random outcome into another random outcome from a
     different distribution.

  2. Inverse

     Useful for "reversing" a transformation to compute one probability in
     terms of another.

  3. `log_det_jacobian(x)`

     "The log of the absolute value of the determinant of the matrix of all
     first-order partial derivatives of the inverse function."

     Useful for inverting a transformation to compute one probability in terms
     of another. Geometrically, the Jacobian determinant is the volume of the
     transformation and is used to scale the probability.

     We take the absolute value of the determinant before log to avoid NaN
     values.  Geometrically, a negative determinant corresponds to an
     orientation-reversing transformation.  It is ok for us to discard the sign
     of the determinant because we only integrate everywhere-nonnegative
     functions (probability densities) and the correct orientation is always the
     one that produces a nonnegative integrand.

  By convention, transformations of random variables are named in terms of the
  forward transformation. The forward transformation creates samples, the
  inverse is useful for computing probabilities.

  #### Example Uses

  - Basic properties:

  ```python
  x = ...  # A tensor.
  # Evaluate forward transformation.
  fwd_x = my_bijector.forward(x)
  x == my_bijector.inverse(fwd_x)
  x != my_bijector.forward(fwd_x)  # Not equal because x != g(g(x)).
  ```

  - Computing a log-likelihood:

  ```python
  def transformed_log_prob(bijector, log_prob, x):
    return (bijector.inverse_log_det_jacobian(x, event_ndims=0) +
            log_prob(bijector.inverse(x)))
  ```

  - Transforming a random outcome:

  ```python
  def transformed_sample(bijector, x):
    return bijector.forward(x)
  ```

  #### Example Bijectors

  - "Exponential"

    ```none
    Y = g(X) = exp(X)
    X ~ Normal(0, 1)  # Univariate.
    ```

    Implies:

    ```none
      g^{-1}(Y) = log(Y)
      |Jacobian(g^{-1})(y)| = 1 / y
      Y ~ LogNormal(0, 1), i.e.,
      prob(Y=y) = |Jacobian(g^{-1})(y)| * prob(X=g^{-1}(y))
                = (1 / y) Normal(log(y); 0, 1)
    ```

    Here is an example of how one might implement the `Exp` bijector:

    ```python
      class Exp(Bijector):

        def __init__(self, validate_args=False, name="exp"):
          super(Exp, self).__init__(
              validate_args=validate_args,
              forward_min_event_ndims=0,
              name=name)

        def _forward(self, x):
          return tf.exp(x)

        def _inverse(self, y):
          return tf.log(y)

        def _inverse_log_det_jacobian(self, y):
          return -self._forward_log_det_jacobian(self._inverse(y))

        def _forward_log_det_jacobian(self, x):
          # Notice that we needn't do any reducing, even when`event_ndims > 0`.
          # The base Bijector class will handle reducing for us; it knows how
          # to do so because we called `super` `__init__` with
          # `forward_min_event_ndims = 0`.
          return x
      ```

  - "Affine"

    ```none
    Y = g(X) = sqrtSigma * X + mu
    X ~ MultivariateNormal(0, I_d)
    ```

    Implies:

    ```none
      g^{-1}(Y) = inv(sqrtSigma) * (Y - mu)
      |Jacobian(g^{-1})(y)| = det(inv(sqrtSigma))
      Y ~ MultivariateNormal(mu, sqrtSigma) , i.e.,
      prob(Y=y) = |Jacobian(g^{-1})(y)| * prob(X=g^{-1}(y))
                = det(sqrtSigma)^(-d) *
                  MultivariateNormal(inv(sqrtSigma) * (y - mu); 0, I_d)
      ```

  #### Min_event_ndims and Naming

  Bijectors are named for the dimensionality of data they act on (i.e. without
  broadcasting). We can think of bijectors having an intrinsic `min_event_ndims`
  , which is the minimum number of dimensions for the bijector act on. For
  instance, a Cholesky decomposition requires a matrix, and hence
  `min_event_ndims=2`.

  Some examples:

  `AffineScalar:  min_event_ndims=0`
  `Affine:  min_event_ndims=1`
  `Cholesky:  min_event_ndims=2`
  `Exp:  min_event_ndims=0`
  `Sigmoid:  min_event_ndims=0`
  `SoftmaxCentered:  min_event_ndims=1`

  Note the difference between `Affine` and `AffineScalar`. `AffineScalar`
  operates on scalar events, whereas `Affine` operates on vector-valued events.

  More generally, there is a `forward_min_event_ndims` and an
  `inverse_min_event_ndims`. In most cases, these will be the same.
  However, for some shape changing bijectors, these will be different
  (e.g. a bijector which pads an extra dimension at the end, might have
  `forward_min_event_ndims=0` and `inverse_min_event_ndims=1`.

  #### Jacobian Determinant

  The Jacobian determinant is a reduction over `event_ndims - min_event_ndims`
  (`forward_min_event_ndims` for `forward_log_det_jacobian` and
  `inverse_min_event_ndims` for `inverse_log_det_jacobian`).
  To see this, consider the `Exp` `Bijector` applied to a `Tensor` which has
  sample, batch, and event (S, B, E) shape semantics. Suppose the `Tensor`'s
  partitioned-shape is `(S=[4], B=[2], E=[3, 3])`. The shape of the `Tensor`
  returned by `forward` and `inverse` is unchanged, i.e., `[4, 2, 3, 3]`.
  However the shape returned by `inverse_log_det_jacobian` is `[4, 2]` because
  the Jacobian determinant is a reduction over the event dimensions.

  Another example is the `Affine` `Bijector`. Because `min_event_ndims = 1`, the
  Jacobian determinant reduction is over `event_ndims - 1`.

  It is sometimes useful to implement the inverse Jacobian determinant as the
  negative forward Jacobian determinant. For example,

  ```python
  def _inverse_log_det_jacobian(self, y):
     return -self._forward_log_det_jac(self._inverse(y))  # Note negation.
  ```

  The correctness of this approach can be seen from the following claim.

  - Claim:

      Assume `Y = g(X)` is a bijection whose derivative exists and is nonzero
      for its domain, i.e., `dY/dX = d/dX g(X) != 0`. Then:

      ```none
      (log o det o jacobian o g^{-1})(Y) = -(log o det o jacobian o g)(X)
      ```

  - Proof:

      From the bijective, nonzero differentiability of `g`, the
      [inverse function theorem](
          https://en.wikipedia.org/wiki/Inverse_function_theorem)
      implies `g^{-1}` is differentiable in the image of `g`.
      Applying the chain rule to `y = g(x) = g(g^{-1}(y))` yields
      `I = g'(g^{-1}(y))*g^{-1}'(y)`.
      The same theorem also implies `g^{-1}'` is non-singular therefore:
      `inv[ g'(g^{-1}(y)) ] = g^{-1}'(y)`.
      The claim follows from [properties of determinant](
  https://en.wikipedia.org/wiki/Determinant#Multiplicativity_and_matrix_groups).

  Generally its preferable to directly implement the inverse Jacobian
  determinant.  This should have superior numerical stability and will often
  share subgraphs with the `_inverse` implementation.

  #### Is_constant_jacobian

  Certain bijectors will have constant jacobian matrices. For instance, the
  `Affine` bijector encodes multiplication by a matrix plus a shift, with
  jacobian matrix, the same aforementioned matrix.

  `is_constant_jacobian` encodes the fact that the jacobian matrix is constant.
  The semantics of this argument are the following:

    * Repeated calls to "log_det_jacobian" functions with the same
      `event_ndims` (but not necessarily same input), will return the first
      computed jacobian (because the matrix is constant, and hence is input
      independent).
    * `log_det_jacobian` implementations are merely broadcastable to the true
      `log_det_jacobian` (because, again, the jacobian matrix is input
      independent). Specifically, `log_det_jacobian` is implemented as the
      log jacobian determinant for a single input.

      ```python
      class Identity(Bijector):

        def __init__(self, validate_args=False, name="identity"):
          super(Identity, self).__init__(
              is_constant_jacobian=True,
              validate_args=validate_args,
              forward_min_event_ndims=0,
              name=name)

        def _forward(self, x):
          return x

        def _inverse(self, y):
          return y

        def _inverse_log_det_jacobian(self, y):
          return -self._forward_log_det_jacobian(self._inverse(y))

        def _forward_log_det_jacobian(self, x):
          # The full log jacobian determinant would be tf.zero_like(x).
          # However, we circumvent materializing that, since the jacobian
          # calculation is input independent, and we specify it for one input.
          return tf.constant(0., x.dtype)

      ```

  #### Subclass Requirements

  - Subclasses typically implement:

      - `_forward`,
      - `_inverse`,
      - `_inverse_log_det_jacobian`,
      - `_forward_log_det_jacobian` (optional).

    The `_forward_log_det_jacobian` is called when the bijector is inverted via
    the `Invert` bijector. If undefined, a slightly less efficiently
    calculation, `-1 * _inverse_log_det_jacobian`, is used.

    If the bijector changes the shape of the input, you must also implement:

      - _forward_event_shape_tensor,
      - _forward_event_shape (optional),
      - _inverse_event_shape_tensor,
      - _inverse_event_shape (optional).

    By default the event-shape is assumed unchanged from input.

  - If the `Bijector`'s use is limited to `TransformedDistribution` (or friends
    like `QuantizedDistribution`) then depending on your use, you may not need
    to implement all of `_forward` and `_inverse` functions.

    Examples:

      1. Sampling (e.g., `sample`) only requires `_forward`.
      2. Probability functions (e.g., `prob`, `cdf`, `survival`) only require
         `_inverse` (and related).
      3. Only calling probability functions on the output of `sample` means
        `_inverse` can be implemented as a cache lookup.

    See "Example Uses" [above] which shows how these functions are used to
    transform a distribution. (Note: `_forward` could theoretically be
    implemented as a cache lookup but this would require controlling the
    underlying sample generation mechanism.)

  #### Non Injective Transforms

  **WARNING** Handing of non-injective transforms is subject to change.

  Non injective maps `g` are supported, provided their domain `D` can be
  partitioned into `k` disjoint subsets, `Union{D1, ..., Dk}`, such that,
  ignoring sets of measure zero, the restriction of `g` to each subset is a
  differentiable bijection onto `g(D)`.  In particular, this imples that for
  `y in g(D)`, the set inverse, i.e. `g^{-1}(y) = {x in D : g(x) = y}`, always
  contains exactly `k` distinct points.

  The property, `_is_injective` is set to `False` to indicate that the bijector
  is not injective, yet satisfies the above condition.

  The usual bijector API is modified in the case `_is_injective is False` (see
  method docstrings for specifics).  Here we show by example the `AbsoluteValue`
  bijector.  In this case, the domain `D = (-inf, inf)`, can be partitioned
  into `D1 = (-inf, 0)`, `D2 = {0}`, and `D3 = (0, inf)`.  Let `gi` be the
  restriction of `g` to `Di`, then both `g1` and `g3` are bijections onto
  `(0, inf)`, with `g1^{-1}(y) = -y`, and `g3^{-1}(y) = y`.  We will use
  `g1` and `g3` to define bijector methods over `D1` and `D3`.  `D2 = {0}` is
  an oddball in that `g2` is one to one, and the derivative is not well defined.
  Fortunately, when considering transformations of probability densities
  (e.g. in `TransformedDistribution`), sets of measure zero have no effect in
  theory, and only a small effect in 32 or 64 bit precision.  For that reason,
  we define `inverse(0)` and `inverse_log_det_jacobian(0)` both as `[0, 0]`,
  which is convenient and results in a left-semicontinuous pdf.

  ```python
  abs = tfp.bijectors.AbsoluteValue()

  abs.forward(-1.)
  ==> 1.

  abs.forward(1.)
  ==> 1.

  abs.inverse(1.)
  ==> (-1., 1.)

  # The |dX/dY| is constant, == 1.  So Log|dX/dY| == 0.
  abs.inverse_log_det_jacobian(1., event_ndims=0)
  ==> (0., 0.)

  # Special case handling of 0.
  abs.inverse(0.)
  ==> (0., 0.)

  abs.inverse_log_det_jacobian(0., event_ndims=0)
  ==> (0., 0.)
  ```

  """

  @abc.abstractmethod
  def __init__(self,
               graph_parents=None,
               is_constant_jacobian=False,
               validate_args=False,
               dtype=None,
               forward_min_event_ndims=None,
               inverse_min_event_ndims=None,
               name=None):
    """Constructs Bijector.

    A `Bijector` transforms random variables into new random variables.

    Examples:

    ```python
    # Create the Y = g(X) = X transform.
    identity = Identity()

    # Create the Y = g(X) = exp(X) transform.
    exp = Exp()
    ```

    See `Bijector` subclass docstring for more details and specific examples.

    Args:
      graph_parents: Python list of graph prerequisites of this `Bijector`.
      is_constant_jacobian: Python `bool` indicating that the Jacobian matrix is
        not a function of the input.
      validate_args: Python `bool`, default `False`. Whether to validate input
        with asserts. If `validate_args` is `False`, and the inputs are invalid,
        correct behavior is not guaranteed.
      dtype: `tf.dtype` supported by this `Bijector`. `None` means dtype is not
        enforced.
      forward_min_event_ndims: Python `integer` indicating the minimum number of
        dimensions `forward` operates on.
      inverse_min_event_ndims: Python `integer` indicating the minimum number of
        dimensions `inverse` operates on. Will be set to
        `forward_min_event_ndims` by default, if no value is provided.
      name: The name to give Ops created by the initializer.

    Raises:
      ValueError:  If neither `forward_min_event_ndims` and
        `inverse_min_event_ndims` are specified, or if either of them is
        negative.
      ValueError:  If a member of `graph_parents` is not a `Tensor`.
    """
    self._graph_parents = graph_parents or []

    if forward_min_event_ndims is None and inverse_min_event_ndims is None:
      raise ValueError("Must specify at least one of `forward_min_event_ndims` "
                       "and `inverse_min_event_ndims`.")
    elif inverse_min_event_ndims is None:
      inverse_min_event_ndims = forward_min_event_ndims
    elif forward_min_event_ndims is None:
      forward_min_event_ndims = inverse_min_event_ndims

    if not isinstance(forward_min_event_ndims, int):
      raise TypeError("Expected forward_min_event_ndims to be of "
                      "type int, got {}".format(
                          type(forward_min_event_ndims).__name__))

    if not isinstance(inverse_min_event_ndims, int):
      raise TypeError("Expected inverse_min_event_ndims to be of "
                      "type int, got {}".format(
                          type(inverse_min_event_ndims).__name__))

    if forward_min_event_ndims < 0:
      raise ValueError("forward_min_event_ndims must be a non-negative "
                       "integer.")
    if inverse_min_event_ndims < 0:
      raise ValueError("inverse_min_event_ndims must be a non-negative "
                       "integer.")

    self._forward_min_event_ndims = forward_min_event_ndims
    self._inverse_min_event_ndims = inverse_min_event_ndims
    self._is_constant_jacobian = is_constant_jacobian
    # Keyed by the current graph.
    self._constant_ildj = {}
    self._validate_args = validate_args
    self._dtype = dtype
    self._from_y = WeakKeyDefaultDict()
    self._from_x = WeakKeyDefaultDict()
    if name:
      self._name = name
    else:
      # We want the default convention to be snake_case rather than CamelCase
      # since `Chain` uses bijector.name as the kwargs dictionary key.
      def camel_to_snake(name):
        s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
        return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()

      self._name = camel_to_snake(type(self).__name__.lstrip("_"))

    for i, t in enumerate(self._graph_parents):
      if t is None or not tf.is_tensor(t):
        raise ValueError("Graph parent item %d is not a Tensor; %s." % (i, t))

  @property
  def graph_parents(self):
    """Returns this `Bijector`'s graph_parents as a Python list."""
    return self._graph_parents

  @property
  def forward_min_event_ndims(self):
    """Returns the minimal number of dimensions bijector.forward operates on."""
    return self._forward_min_event_ndims

  @property
  def inverse_min_event_ndims(self):
    """Returns the minimal number of dimensions bijector.inverse operates on."""
    return self._inverse_min_event_ndims

  @property
  def is_constant_jacobian(self):
    """Returns true iff the Jacobian matrix is not a function of x.

    Note: Jacobian matrix is either constant for both forward and inverse or
    neither.

    Returns:
      is_constant_jacobian: Python `bool`.
    """
    return self._is_constant_jacobian

  @property
  def _is_injective(self):
    """Returns true iff the forward map `g` is injective (one-to-one function).

    **WARNING** This hidden property and its behavior are subject to change.

    Note:  Non-injective maps `g` are supported, provided their domain `D` can
    be partitioned into `k` disjoint subsets, `Union{D1, ..., Dk}`, such that,
    ignoring sets of measure zero, the restriction of `g` to each subset is a
    differentiable bijection onto `g(D)`.

    Returns:
      is_injective: Python `bool`.
    """
    return True

  @property
  def validate_args(self):
    """Returns True if Tensor arguments will be validated."""
    return self._validate_args

  @property
  def dtype(self):
    """dtype of `Tensor`s transformable by this distribution."""
    return self._dtype

  @property
  def name(self):
    """Returns the string name of this `Bijector`."""
    return self._name

  def __call__(self, value, name=None, **kwargs):
    """Applies or composes the `Bijector`, depending on input type.

    This is a convenience function which applies the `Bijector` instance in
    three different ways, depending on the input:

    1. If the input is a `tfd.Distribution` instance, return
       `tfd.TransformedDistribution(distribution=input, bijector=self)`.
    2. If the input is a `tfb.Bijector` instance, return
       `tfb.Chain([self, input])`.
    3. Otherwise, return `self.forward(input)`

    Args:
      value: A `tfd.Distribution`, `tfb.Bijector`, or a `Tensor`.
      name: Python `str` name given to ops created by this function.
      **kwargs: Additional keyword arguments passed into the created
        `tfd.TransformedDistribution`, `tfb.Bijector`, or `self.forward`.

    Returns:
      composition: A `tfd.TransformedDistribution` if the input was a
        `tfd.Distribution`, a `tfb.Chain` if the input was a `tfb.Bijector`, or
        a `Tensor` computed by `self.forward`.

    #### Examples

    ```python
    sigmoid = tfb.Reciprocal()(
        tfb.AffineScalar(shift=1.)(
          tfb.Exp()(
            tfb.AffineScalar(scale=-1.))))
    # ==> `tfb.Chain([
    #         tfb.Reciprocal(),
    #         tfb.AffineScalar(shift=1.),
    #         tfb.Exp(),
    #         tfb.AffineScalar(scale=-1.),
    #      ])`  # ie, `tfb.Sigmoid()`

    log_normal = tfb.Exp()(tfd.Normal(0, 1))
    # ==> `tfd.TransformedDistribution(tfd.Normal(0, 1), tfb.Exp())`

    tfb.Exp()([-1., 0., 1.])
    # ==> tf.exp([-1., 0., 1.])
    ```

    """

    # To avoid circular dependencies and keep the implementation local to the
    # `Bijector` class, we violate PEP8 guidelines and import here rather than
    # at the top of the file.
    from tensorflow_probability.python.bijectors import chain  # pylint: disable=g-import-not-at-top
    from tensorflow_probability.python.distributions import distribution  # pylint: disable=g-import-not-at-top
    from tensorflow_probability.python.distributions import transformed_distribution  # pylint: disable=g-import-not-at-top

    # TODO(b/128841942): Handle Conditional distributions and bijectors.
    if type(value) is transformed_distribution.TransformedDistribution:  # pylint: disable=unidiomatic-typecheck
      # We cannot accept subclasses with different constructors here, because
      # subclass constructors may accept constructor arguments TD doesn't know
      # how to handle. e.g. `TypeError: __init__() got an unexpected keyword
      # argument 'allow_nan_stats'` when doing
      # `tfb.Identity()(tfd.Chi(df=1., allow_nan_stats=True))`.
      new_kwargs = value.parameters
      new_kwargs.update(kwargs)
      new_kwargs["name"] = name or new_kwargs.get("name", None)
      new_kwargs["bijector"] = self(value.bijector)
      return transformed_distribution.TransformedDistribution(**new_kwargs)

    if isinstance(value, distribution.Distribution):
      return transformed_distribution.TransformedDistribution(
          distribution=value,
          bijector=self,
          name=name,
          **kwargs)

    if isinstance(value, chain.Chain):
      new_kwargs = kwargs.copy()
      new_kwargs["bijectors"] = [self] + ([] if value.bijectors is None
                                          else list(value.bijectors))
      if "validate_args" not in new_kwargs:
        new_kwargs["validate_args"] = value.validate_args
      new_kwargs["name"] = name or value.name
      return chain.Chain(**new_kwargs)

    if isinstance(value, Bijector):
      return chain.Chain([self, value], name=name, **kwargs)

    return self._call_forward(value, name=name or "forward", **kwargs)

  def _forward_event_shape_tensor(self, input_shape):
    """Subclass implementation for `forward_event_shape_tensor` function."""
    # By default, we assume event_shape is unchanged.
    return input_shape

  def forward_event_shape_tensor(self,
                                 input_shape,
                                 name="forward_event_shape_tensor"):
    """Shape of a single sample from a single batch as an `int32` 1D `Tensor`.

    Args:
      input_shape: `Tensor`, `int32` vector indicating event-portion shape
        passed into `forward` function.
      name: name to give to the op

    Returns:
      forward_event_shape_tensor: `Tensor`, `int32` vector indicating
        event-portion shape after applying `forward`.
    """
    with self._name_scope(name):
      input_shape = tf.convert_to_tensor(
          value=input_shape, dtype=tf.int32, name="input_shape")
      return self._forward_event_shape_tensor(input_shape)

  def _forward_event_shape(self, input_shape):
    """Subclass implementation for `forward_event_shape` public function."""
    # By default, we assume event_shape is unchanged.
    return input_shape

  def forward_event_shape(self, input_shape):
    """Shape of a single sample from a single batch as a `TensorShape`.

    Same meaning as `forward_event_shape_tensor`. May be only partially defined.

    Args:
      input_shape: `TensorShape` indicating event-portion shape passed into
        `forward` function.

    Returns:
      forward_event_shape_tensor: `TensorShape` indicating event-portion shape
        after applying `forward`. Possibly unknown.
    """
    return self._forward_event_shape(tf.TensorShape(input_shape))

  def _inverse_event_shape_tensor(self, output_shape):
    """Subclass implementation for `inverse_event_shape_tensor` function."""
    # By default, we assume event_shape is unchanged.
    return output_shape

  def inverse_event_shape_tensor(self,
                                 output_shape,
                                 name="inverse_event_shape_tensor"):
    """Shape of a single sample from a single batch as an `int32` 1D `Tensor`.

    Args:
      output_shape: `Tensor`, `int32` vector indicating event-portion shape
        passed into `inverse` function.
      name: name to give to the op

    Returns:
      inverse_event_shape_tensor: `Tensor`, `int32` vector indicating
        event-portion shape after applying `inverse`.
    """
    with self._name_scope(name):
      output_shape = tf.convert_to_tensor(
          value=output_shape, dtype=tf.int32, name="output_shape")
      return self._inverse_event_shape_tensor(output_shape)

  def _inverse_event_shape(self, output_shape):
    """Subclass implementation for `inverse_event_shape` public function."""
    # By default, we assume event_shape is unchanged.
    return tf.TensorShape(output_shape)

  def inverse_event_shape(self, output_shape):
    """Shape of a single sample from a single batch as a `TensorShape`.

    Same meaning as `inverse_event_shape_tensor`. May be only partially defined.

    Args:
      output_shape: `TensorShape` indicating event-portion shape passed into
        `inverse` function.

    Returns:
      inverse_event_shape_tensor: `TensorShape` indicating event-portion shape
        after applying `inverse`. Possibly unknown.
    """
    return self._inverse_event_shape(tf.TensorShape(output_shape))

  def _forward(self, x):
    """Subclass implementation for `forward` public function."""
    raise NotImplementedError("forward not implemented.")

  def _call_forward(self, x, name, **kwargs):
    """Wraps call to _forward, allowing extra shared logic."""
    with self._name_scope(name):
      x = tf.convert_to_tensor(value=x, name="x")
      self._maybe_assert_dtype(x)
      if not self._is_injective:  # No caching for non-injective
        return self._forward(x, **kwargs)
      mapping = self._lookup(x=x, kwargs=kwargs)
      if mapping.y is not None:
        return mapping.y
      mapping = mapping.merge(y=self._forward(x, **kwargs))
      # It's most important to cache the y->x mapping, because computing
      # inverse(forward(y)) may be numerically unstable / lossy. Caching the
      # x->y mapping only saves work. Since python doesn't support ephemerons,
      # we cannot be simultaneously weak-keyed on both x and y, so we choose y.
      self._cache_by_y(mapping)
      if not tf.executing_eagerly():
        self._cache_by_x(mapping)
      return mapping.y

  def forward(self, x, name="forward", **kwargs):
    """Returns the forward `Bijector` evaluation, i.e., X = g(Y).

    Args:
      x: `Tensor`. The input to the "forward" evaluation.
      name: The name to give this op.
      **kwargs: Named arguments forwarded to subclass implementation.

    Returns:
      `Tensor`.

    Raises:
      TypeError: if `self.dtype` is specified and `x.dtype` is not
        `self.dtype`.
      NotImplementedError: if `_forward` is not implemented.
    """
    return self._call_forward(x, name, **kwargs)

  def _inverse(self, y):
    """Subclass implementation for `inverse` public function."""
    raise NotImplementedError("inverse not implemented")

  def _call_inverse(self, y, name, **kwargs):
    """Wraps call to _inverse, allowing extra shared logic."""
    with self._name_scope(name):
      y = tf.convert_to_tensor(value=y, name="y")
      self._maybe_assert_dtype(y)
      if not self._is_injective:  # No caching for non-injective
        return self._inverse(y, **kwargs)
      mapping = self._lookup(y=y, kwargs=kwargs)
      if mapping.x is not None:
        return mapping.x
      mapping = mapping.merge(x=self._inverse(y, **kwargs))
      # It's most important to cache the x->y mapping, because computing
      # forward(inverse(y)) may be numerically unstable / lossy. Caching the
      # y->x mapping only saves work. Since python doesn't support ephemerons,
      # we cannot be simultaneously weak-keyed on both x and y, so we choose x.
      self._cache_by_x(mapping)
      if not tf.executing_eagerly():
        self._cache_by_y(mapping)
      return mapping.x

  def inverse(self, y, name="inverse", **kwargs):
    """Returns the inverse `Bijector` evaluation, i.e., X = g^{-1}(Y).

    Args:
      y: `Tensor`. The input to the "inverse" evaluation.
      name: The name to give this op.
      **kwargs: Named arguments forwarded to subclass implementation.

    Returns:
      `Tensor`, if this bijector is injective.
        If not injective, returns the k-tuple containing the unique
        `k` points `(x1, ..., xk)` such that `g(xi) = y`.

    Raises:
      TypeError: if `self.dtype` is specified and `y.dtype` is not
        `self.dtype`.
      NotImplementedError: if `_inverse` is not implemented.
    """
    return self._call_inverse(y, name, **kwargs)

  def _compute_inverse_log_det_jacobian_with_caching(
      self, x, y, prefer_inverse_ldj_fn, event_ndims, kwargs):
    """Compute ILDJ by the best available means, and ensure it is cached.

    Sub-classes of Bijector may implement either `_forward_log_det_jacobian` or
    `_inverse_log_det_jacobian`, or both, and may prefer one or the other for
    reasons of computational efficiency and/or numerics. In general, to compute
    the [I]LDJ, we need one of `x` or `y` and one of `_forward_log_det_jacobian`
    or `_inverse_log_det_jacobian` (all bijectors implement `forward` and
    `inverse`, so either of `x` and `y` may always computed from the other).

    This method encapsulates the logic of selecting the best possible method of
    computing the inverse log det jacobian, and executing it. Possible avenues
    to obtaining this value are:
      - recovering it from the cache,
      - computing it as `-_forward_log_det_jacobian(x)`, where `x` may need to
        be computed as `inverse(y)`, or
      - computing it as `_inverse_log_det_jacobian(y)`, where `y` may need to
        be computed as `forward(x)`.

    To make things more interesting, we may be able to take advantage of the
    knowledge that the jacobian is constant, for given `event_ndims`, and may
    even cache that result.

    To make things even more interesting still, we may need to perform an
    additional `reduce_sum` over the event dims of an input after computing the
    ILDJ (see `_reduce_jacobian_det_over_event` for the reduction logic). In
    order to prevent spurious TF graph dependencies on past inputs in cached
    results, we need to take care to do this reduction after the cache lookups.

    This method takes care of all the above concerns.

    Args:
      x: a `Tensor`, the pre-Bijector transform value at whose post-Bijector
        transform value the ILDJ is to be computed. Can be `None` as long as
        `y` is not `None`.
      y: a `Tensor`, a point in the output space of the Bijector's `forward`
        transformation, at whose value the ILDJ is to be computed. Can be
        `None` as long as `x` is not `None`.
      prefer_inverse_ldj_fn: Python `bool`, if `True`, will strictly prefer to
        use the `_inverse_log_det_jacobian` to compute ILDJ; else, will
        strictly prefer to use `_forward_log_det_jacobian`.
      event_ndims: int-like `Tensor`, the number of dims of an event (in the
        pre- or post-transformed space, as appropriate). These need to be summed
        over to compute the total ildj.
      kwargs: dictionary of keyword args that will be passed to calls to
        `self.forward` or `self.inverse` (if either of those are required), and
        to `self._compute_unreduced_nonconstant_ildj_with_caching` or
        `self._compute_unreduced_constant_ildj_with_caching`, as appropriate (
        those functions use the conditioning kwargs for caching and for
        their underlying computations).

    Returns:
      ildj: a Tensor of ILDJ['s] at the given input (as specified by the args).
        Also updates the cache as needed.
    """
    # Ensure at least one of _inverse/_forward_log_det_jacobian is defined.
    if not (hasattr(self, "_inverse_log_det_jacobian") or
            hasattr(self, "_forward_log_det_jacobian")):
      raise NotImplementedError(
          "Neither _forward_log_det_jacobian nor _inverse_log_det_jacobian "
          "is implemented. One or the other is required.")

    # Use inverse_log_det_jacobian if either
    #   1. it is preferred to *and* we are able, or
    #   2. forward ldj fn isn't implemented (so we have no choice).
    use_inverse_ldj_fn = (
        (prefer_inverse_ldj_fn and hasattr(self, "_inverse_log_det_jacobian"))
        or not hasattr(self, "_forward_log_det_jacobian"))

    if use_inverse_ldj_fn:
      tensor_to_use = y if y is not None else self.forward(x, **kwargs)
      min_event_ndims = self.inverse_min_event_ndims
    else:
      tensor_to_use = x if x is not None else self.inverse(y, **kwargs)
      min_event_ndims = self.forward_min_event_ndims

    if self.is_constant_jacobian:
      unreduced_ildj = self._compute_unreduced_constant_ildj_with_caching(
          tensor_to_use, use_inverse_ldj_fn, kwargs)
    else:
      unreduced_ildj = self._compute_unreduced_nonconstant_ildj_with_caching(
          x, y, tensor_to_use, use_inverse_ldj_fn, kwargs)

    return self._reduce_jacobian_det_over_event(
        tf.shape(input=tensor_to_use), unreduced_ildj, min_event_ndims,
        event_ndims)

  def _compute_unreduced_nonconstant_ildj_with_caching(
      self, x, y, tensor_to_use, use_inverse_ldj_fn, kwargs):
    """Helper for computing ILDJ, with caching, in the non-constant case.

    Does not do the "reduce" step which is necessary in some cases; this is left
    to the caller.

    Args:
      x: a `Tensor`, the pre-Bijector transform value at whose post-Bijector
        transform value the ILDJ is to be computed. Can be `None` as long as
        `y` is not `None`. This method only uses the value for cache
        lookup/updating.
      y: a `Tensor`, a point in the output space of the Bijector's `forward`
        transformation, at whose value the ILDJ is to be computed. Can be
        `None` as long as `x` is not `None`. This method only uses the value
        for cache lookup/updating.
      tensor_to_use: a `Tensor`, the one to actually pass to the chosen compute
        function (`_inverse_log_det_jacobian` or `_forward_log_det_jacobian`).
        It is presumed that the caller has already figured out what input to use
        (it will either be the x or y value corresponding to the location where
        we are computing the ILDJ).
      use_inverse_ldj_fn: Python `bool`, if `True`, will use the
        `_inverse_log_det_jacobian` to compute ILDJ; else, will use
        `_forward_log_det_jacobian`.
      kwargs: dictionary of keyword args that will be passed to calls to to
        `_inverse_log_det_jacobian` or `_forward_log_det_jacobian`, as well as
        for lookup/updating of the result in the cache.

    Returns:
      ildj: the (un-reduce_sum'ed) value of the ILDJ at the specified input
      location. Also updates the cache as needed.
    """
    mapping = self._lookup(x=x, y=y, kwargs=kwargs)
    if mapping.ildj is not None:
      return mapping.ildj

    if use_inverse_ldj_fn:
      ildj = self._inverse_log_det_jacobian(tensor_to_use, **kwargs)
    else:
      ildj = -self._forward_log_det_jacobian(tensor_to_use, **kwargs)

    mapping = mapping.merge(x=x, y=y, ildj=ildj)
    self._cache_update(mapping)
    return ildj

  def _compute_unreduced_constant_ildj_with_caching(
      self, tensor_to_use, use_inverse_ldj_fn, kwargs):
    """Helper for computing ILDJ, with caching, in the constant-ILDJ case.

    Does not do the "reduce" step which is necessary in some cases; this is left
    to the caller.

    Args:
      tensor_to_use: a `Tensor`, to pass to the chosen compute
        function (`_inverse_log_det_jacobian` or `_forward_log_det_jacobian`).
        Since we know the ILDJ is input-independent, the actual value is not
        meaningful, although the shape and dtype may be.
      use_inverse_ldj_fn: Python `bool`, if `True`, will use the
        `_inverse_log_det_jacobian` to compute ILDJ; else, will use
        `_forward_log_det_jacobian`.
      kwargs: dictionary of keyword args that will be passed to calls to
        to `_inverse_log_det_jacobian` or `_forward_log_det_jacobian`, as well
        as for lookup/updating of the result in the cache.

    Returns:
      ildj: the (un-reduce_sum'ed) value of the ILDJ for the specified input.
        Also updates the cache as needed.
    """
    current_graph = _get_current_graph()
    if current_graph in self._constant_ildj:
      return self._constant_ildj[current_graph]

    if use_inverse_ldj_fn:
      ildj = self._inverse_log_det_jacobian(tensor_to_use, **kwargs)
    else:
      ildj = -self._forward_log_det_jacobian(tensor_to_use, **kwargs)
    self._constant_ildj[current_graph] = ildj
    return ildj

  def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs):
    """Wraps call to _inverse_log_det_jacobian, allowing extra shared logic.

    Specifically, this method
      - adds a name scope,
      - performs validations,
      - handles the special case of non-injective Bijector (skip caching and
        reduce_sum over the values at the multiple points in the preimage of `y`
        under the non-injective transformation)

    so that sub-classes don't have to worry about this stuff.

    Args:
      y: same as in `inverse_log_det_jacobian`
      event_ndims: same as in `inverse_log_det_jacobian`
      name: same as in `inverse_log_det_jacobian`
      **kwargs: same as in `inverse_log_det_jacobian`

    Returns:
      ildj: the inverse log det jacobian at `y`. Also updates the cache as
        needed.
    """
    with self._name_scope(name), tf.control_dependencies(
        self._check_valid_event_ndims(
            min_event_ndims=self.inverse_min_event_ndims,
            event_ndims=event_ndims)):
      y = tf.convert_to_tensor(value=y, name="y")
      self._maybe_assert_dtype(y)

      if not self._is_injective:
        ildjs = self._inverse_log_det_jacobian(y, **kwargs)
        return tuple(
            self._reduce_jacobian_det_over_event(  # pylint: disable=g-complex-comprehension
                tf.shape(input=y),
                ildj,
                self.inverse_min_event_ndims, event_ndims)
            for ildj in ildjs)

      return self._compute_inverse_log_det_jacobian_with_caching(
          x=None,
          y=y,
          prefer_inverse_ldj_fn=True,
          event_ndims=event_ndims,
          kwargs=kwargs)

  def inverse_log_det_jacobian(self,
                               y,
                               event_ndims,
                               name="inverse_log_det_jacobian",
                               **kwargs):
    """Returns the (log o det o Jacobian o inverse)(y).

    Mathematically, returns: `log(det(dX/dY))(Y)`. (Recall that: `X=g^{-1}(Y)`.)

    Note that `forward_log_det_jacobian` is the negative of this function,
    evaluated at `g^{-1}(y)`.

    Args:
      y: `Tensor`. The input to the "inverse" Jacobian determinant evaluation.
      event_ndims: Number of dimensions in the probabilistic events being
        transformed. Must be greater than or equal to
        `self.inverse_min_event_ndims`. The result is summed over the final
        dimensions to produce a scalar Jacobian determinant for each event, i.e.
        it has shape `rank(y) - event_ndims` dimensions.
      name: The name to give this op.
      **kwargs: Named arguments forwarded to subclass implementation.

    Returns:
      ildj: `Tensor`, if this bijector is injective.
        If not injective, returns the tuple of local log det
        Jacobians, `log(det(Dg_i^{-1}(y)))`, where `g_i` is the restriction
        of `g` to the `ith` partition `Di`.

    Raises:
      TypeError: if `self.dtype` is specified and `y.dtype` is not
        `self.dtype`.
      NotImplementedError: if `_inverse_log_det_jacobian` is not implemented.
    """
    return self._call_inverse_log_det_jacobian(y, event_ndims, name, **kwargs)

  def _call_forward_log_det_jacobian(self, x, event_ndims, name, **kwargs):
    """Wraps call to _forward_log_det_jacobian, allowing extra shared logic.

    Specifically, this method
      - adds a name scope,
      - performs validations,
      - handles the special case of non-injective Bijector (forward jacobian is
        ill-defined in this case and we raise an exception)

    so that sub-classes don't have to worry about this stuff.

    Args:
      x: same as in `forward_log_det_jacobian`
      event_ndims: same as in `forward_log_det_jacobian`
      name: same as in `forward_log_det_jacobian`
      **kwargs: same as in `forward_log_det_jacobian`

    Returns:
      fldj: the forward log det jacobian at `x`. Also updates the cache as
      needed.
    """
    if not self._is_injective:
      raise NotImplementedError(
          "forward_log_det_jacobian cannot be implemented for non-injective "
          "transforms.")

    with self._name_scope(name), tf.control_dependencies(
        self._check_valid_event_ndims(
            min_event_ndims=self.forward_min_event_ndims,
            event_ndims=event_ndims)):
      x = tf.convert_to_tensor(value=x, name="x")
      self._maybe_assert_dtype(x)

      return -self._compute_inverse_log_det_jacobian_with_caching(
          x=x,
          y=None,
          prefer_inverse_ldj_fn=False,
          event_ndims=event_ndims,
          kwargs=kwargs)

  def forward_log_det_jacobian(self,
                               x,
                               event_ndims,
                               name="forward_log_det_jacobian",
                               **kwargs):
    """Returns both the forward_log_det_jacobian.

    Args:
      x: `Tensor`. The input to the "forward" Jacobian determinant evaluation.
      event_ndims: Number of dimensions in the probabilistic events being
        transformed. Must be greater than or equal to
        `self.forward_min_event_ndims`. The result is summed over the final
        dimensions to produce a scalar Jacobian determinant for each event, i.e.
        it has shape `rank(x) - event_ndims` dimensions.
      name: The name to give this op.
      **kwargs: Named arguments forwarded to subclass implementation.

    Returns:
      `Tensor`, if this bijector is injective.
        If not injective this is not implemented.

    Raises:
      TypeError: if `self.dtype` is specified and `y.dtype` is not
        `self.dtype`.
      NotImplementedError: if neither `_forward_log_det_jacobian`
        nor {`_inverse`, `_inverse_log_det_jacobian`} are implemented, or
        this is a non-injective bijector.
    """
    return self._call_forward_log_det_jacobian(x, event_ndims, name, **kwargs)

  @contextlib.contextmanager
  def _name_scope(self, name=None):
    """Helper function to standardize op scope."""
    with tf.compat.v2.name_scope(self.name):
      with tf.compat.v2.name_scope(name) as scope:
        yield scope

  def _maybe_assert_dtype(self, x):
    """Helper to check dtype when self.dtype is known."""
    if (self.dtype is not None and
        not dtype_util.base_equal(self.dtype, x.dtype)):
      raise TypeError(
          "Input had dtype %s but expected %s." % (x.dtype, self.dtype))

  def _cache_by_x(self, mapping):
    """Helper which stores new mapping info in the forward dict."""
    # Merging from lookup is an added check that we're not overwriting anything
    # which is not None.
    mapping = mapping.merge(
        mapping=self._lookup(mapping.x, mapping.y, mapping.kwargs))
    if mapping.x is None:
      raise ValueError("Caching expects x to be known, i.e., not None.")
    self._from_x[mapping.x][mapping.subkey] = mapping.remove("x")

  def _cache_by_y(self, mapping):
    """Helper which stores new mapping info in the inverse dict."""
    # Merging from lookup is an added check that we're not overwriting anything
    # which is not None.
    mapping = mapping.merge(
        mapping=self._lookup(mapping.x, mapping.y, mapping.kwargs))
    if mapping.y is None:
      raise ValueError("Caching expects y to be known, i.e., not None.")
    self._from_y[mapping.y][mapping.subkey] = mapping.remove("y")

  def _cache_update(self, mapping):
    """Helper which updates only those cached entries that already exist."""
    if mapping.x is not None and mapping.subkey in self._from_x[mapping.x]:
      self._cache_by_x(mapping)
    if mapping.y is not None and mapping.subkey in self._from_y[mapping.y]:
      self._cache_by_y(mapping)

  def _lookup(self, x=None, y=None, kwargs=None):
    """Helper which retrieves mapping info from forward/inverse dicts."""
    mapping = _Mapping(x=x, y=y, kwargs=kwargs)
    subkey = mapping.subkey
    if x is not None:
      # We removed x at caching time. Add it back if we lookup successfully.
      mapping = self._from_x[x].get(subkey, mapping).merge(x=x)
    if y is not None:
      # We removed y at caching time. Add it back if we lookup successfully.
      mapping = self._from_y[y].get(subkey, mapping).merge(y=y)
    return mapping

  def _reduce_jacobian_det_over_event(
      self, shape_tensor, ildj, min_event_ndims, event_ndims):
    """Reduce jacobian over event_ndims - min_event_ndims."""
    # In this case, we need to tile the Jacobian over the event and reduce.
    rank = tf.size(input=shape_tensor)
    shape_tensor = shape_tensor[rank - event_ndims:rank - min_event_ndims]

    ones = tf.ones(shape_tensor, ildj.dtype)
    reduced_ildj = tf.reduce_sum(
        input_tensor=ones * ildj,
        axis=self._get_event_reduce_dims(min_event_ndims, event_ndims))

    return reduced_ildj

  def _get_event_reduce_dims(self, min_event_ndims, event_ndims):
    """Compute the reduction dimensions given event_ndims."""
    event_ndims_ = self._maybe_get_static_event_ndims(event_ndims)

    if event_ndims_ is not None:
      return [-index for index in range(1, event_ndims_ - min_event_ndims + 1)]
    else:
      reduce_ndims = event_ndims - min_event_ndims
      return tf.range(-reduce_ndims, 0)

  def _check_valid_event_ndims(self, min_event_ndims, event_ndims):
    """Check whether event_ndims is atleast min_event_ndims."""
    event_ndims = tf.convert_to_tensor(value=event_ndims, name="event_ndims")
    event_ndims_ = tf.get_static_value(event_ndims)
    assertions = []

    if not dtype_util.is_integer(event_ndims.dtype):
      raise ValueError("Expected integer dtype, got dtype {}".format(
          event_ndims.dtype))

    if event_ndims_ is not None:
      if tensorshape_util.rank(event_ndims.shape) != 0:
        raise ValueError("Expected scalar event_ndims, got shape {}".format(
            event_ndims.shape))
      if min_event_ndims > event_ndims_:
        raise ValueError("event_ndims ({}) must be larger than "
                         "min_event_ndims ({})".format(event_ndims_,
                                                       min_event_ndims))
    elif self.validate_args:
      assertions += [
          assert_util.assert_greater_equal(event_ndims, min_event_ndims)
      ]

    if tensorshape_util.is_fully_defined(event_ndims.shape):
      if tensorshape_util.rank(event_ndims.shape) != 0:
        raise ValueError("Expected scalar shape, got ndims {}".format(
            tensorshape_util.rank(event_ndims.shape)))

    elif self.validate_args:
      assertions += [
          assert_util.assert_rank(event_ndims, 0, message="Expected scalar.")
      ]
    return assertions

  def _maybe_get_static_event_ndims(self, event_ndims):
    """Helper which returns tries to return an integer static value."""
    event_ndims_ = distribution_util.maybe_get_static_value(event_ndims)

    if isinstance(event_ndims_, (np.generic, np.ndarray)):
      if event_ndims_.dtype not in (np.int32, np.int64):
        raise ValueError("Expected integer dtype, got dtype {}".format(
            event_ndims_.dtype))

      if isinstance(event_ndims_, np.ndarray) and len(event_ndims_.shape):
        raise ValueError(
            "Expected a scalar integer, got {}".format(event_ndims_))
      event_ndims_ = int(event_ndims_)

    return event_ndims_

class ConditionalBijector(Bijector):
  """Conditional Bijector is a Bijector that allows intrinsic conditioning."""

  @deprecation.deprecated(
      "2019-07-01",
      "`ConditionalBijector` is no longer required; `Bijector` "
      "top-level functions now pass-through `**kwargs`.",
      warn_once=True)
  def __new__(cls, *args, **kwargs):  # pylint: disable=unused-argument
    return super(ConditionalBijector, cls).__new__(cls)


# 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 to apply slice sampling update in one dimension."""

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

import tensorflow as tf

from tensorflow_probability.python import distributions

def _left_doubling_increments(batch_shape, max_doublings, step_size, seed=None,
                              name=None):
  """Computes the doubling increments for the left end point.

  The doubling procedure expands an initial interval to find a superset of the
  true slice. At each doubling iteration, the interval width is doubled to
  either the left or the right hand side with equal probability.
  If, initially, the left end point is at `L(0)` and the width of the
  interval is `w(0)`, then the left end point and the width at the
  k-th iteration (denoted L(k) and w(k) respectively) are given by the following
  recursions:

  ```none
  w(k) = 2 * w(k-1)
  L(k) = L(k-1) - w(k-1) * X_k, X_k ~ Bernoulli(0.5)
  or, L(0) - L(k) = w(0) Sum(2^i * X(i+1), 0 <= i < k)
  ```

  This function computes the sequence of `L(0)-L(k)` and `w(k)` for k between 0
  and `max_doublings` independently for each chain.

  Args:
    batch_shape: Positive int32 `tf.Tensor`. The batch shape.
    max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
      doublings to consider.
    step_size: A real `tf.Tensor` with shape compatible with [num_chains].
      The size of the initial interval.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    left_increments: A tensor of shape (max_doublings+1, batch_shape). The
      relative position of the left end point after the doublings.
    widths: A tensor of shape (max_doublings+1, ones_like(batch_shape)). The
      widths of the intervals at each stage of the doubling.
  """
  with tf.compat.v1.name_scope(name, 'left_doubling_increments',
                               [batch_shape, max_doublings, step_size]):

    step_size = tf.convert_to_tensor(value=step_size)
    dtype = step_size.dtype.base_dtype
    # Output shape of the left increments tensor.
    output_shape = tf.concat(([max_doublings + 1], batch_shape), axis=0)
    # A sample realization of X_k.
    expand_left = distributions.Bernoulli(0.5, dtype=dtype).sample(
        sample_shape=output_shape, seed=seed)

    # The widths of the successive intervals. Starts with 1.0 and ends with
    # 2^max_doublings.
    width_multipliers = tf.cast(2 ** tf.range(0, max_doublings+1), dtype=dtype)
    # Output shape of the `widths` tensor.
    widths_shape = tf.concat(([max_doublings + 1],
                              tf.ones_like(batch_shape)), axis=0)
    width_multipliers = tf.reshape(width_multipliers, shape=widths_shape)
    # Widths shape is [max_doublings + 1, 1, 1, 1...].
    widths = width_multipliers * step_size

    # Take the cumulative sum of the left side increments in slice width to give
    # the resulting distance from the inital lower bound.
    left_increments = tf.cumsum(widths * expand_left, exclusive=True, axis=0)
    return left_increments, widths

def _find_best_interval_idx(x, name=None):
  """Finds the index of the optimal set of bounds for each chain.

  For each chain, finds the smallest set of bounds for which both edges lie
  outside the slice. This is equivalent to the point at which a for loop
  implementation (P715 of Neal (2003)) of the algorithm would terminate.

  Performs the following calculation, where i is the number of doublings that
  have been performed and k is the max number of doublings:

  (2 * k - i) * flag + i

  The argmax of the above returns the earliest index where the bounds were
  outside the slice and if there is no such point, the widest bounds.

  Args:
    x: A tensor of shape (max_doublings+1, batch_shape). Type int32, with value
      0 or 1. Indicates if this set of bounds is outside the slice.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    indices: A tensor of shape batch_shape. Type int32, with the index of the
      first set of bounds outside the slice and if there are none, the index of
      the widest set.
  """
  with tf.compat.v1.name_scope(name, 'find_best_interval_idx', [x]):
    # Returns max_doublings + 1. Positive int32.
    k = tf.shape(input=x)[0]
    dtype = x.dtype.base_dtype
    # Factors by which to multiply the flag. Corresponds to (2 * k - i) above.
    mults = tf.range(2 * k, k, -1, dtype=dtype)[:, tf.newaxis]
    # Factors by which to shift the flag. Corresponds to i above. Ensures the
    # widest bounds are selected if there are no bounds outside the slice.
    shifts = tf.range(k, dtype=dtype)[:, tf.newaxis]
    indices = tf.argmax(input=mults * x + shifts, axis=0, output_type=dtype)
    return indices

def slice_bounds_by_doubling(x_initial,
                             target_log_prob,
                             log_slice_heights,
                             max_doublings,
                             step_size,
                             seed=None,
                             name=None):
  """Returns the bounds of the slice at each stage of doubling procedure.

  Precomputes the x coordinates of the left (L) and right (R) endpoints of the
  interval `I` produced in the "doubling" algorithm [Neal 2003][1] P713. Note
  that we simultaneously compute all possible doubling values for each chain,
  for the reason that at small-medium densities, the gains from parallel
  evaluation might cause a speed-up, but this will be benchmarked against the
  while loop implementation.

  Args:
    x_initial: `tf.Tensor` of any shape and any real dtype consumable by
      `target_log_prob`. The initial points.
    target_log_prob: A callable taking a `tf.Tensor` of shape and dtype as
      `x_initial` and returning a tensor of the same shape. The log density of
      the target distribution.
    log_slice_heights: `tf.Tensor` with the same shape as `x_initial` and the
      same dtype as returned by `target_log_prob`. The log of the height of the
      slice for each chain. The values must be bounded above by
      `target_log_prob(x_initial)`.
    max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
      doublings to consider.
    step_size: `tf.Tensor` with same dtype as and shape compatible with
      `x_initial`. The size of the initial interval.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    upper_bounds: A tensor of same shape and dtype as `x_initial`. Slice upper
      bounds for each chain.
    lower_bounds: A tensor of same shape and dtype as `x_initial`. Slice lower
      bounds for each chain.
    both_ok: A tensor of shape `x_initial` and boolean dtype. Indicates if both
      the chosen upper and lower bound lie outside of the slice.

  #### References

  [1]: Radford M. Neal. Slice Sampling. The Annals of Statistics. 2003, Vol 31,
       No. 3 , 705-767.
       https://projecteuclid.org/download/pdf_1/euclid.aos/1056562461
  """
  with tf.compat.v1.name_scope(
      name, 'slice_bounds_by_doubling',
      [x_initial, log_slice_heights, max_doublings, step_size]):
    seed_gen = distributions.SeedStream(seed, salt='slice_bounds_by_doubling')
    x_initial = tf.convert_to_tensor(value=x_initial)
    batch_shape = tf.shape(input=x_initial)
    dtype = step_size.dtype.base_dtype
    left_endpoints = x_initial + step_size * tf.random.uniform(
        batch_shape, minval=-1.0, maxval=0.0, dtype=dtype, seed=seed_gen())

    # Compute the increments by which we need to step the upper and lower bounds
    # part of the doubling procedure.
    left_increments, widths = _left_doubling_increments(
        batch_shape, max_doublings, step_size, seed=seed_gen())
    # The left and right end points. Shape (max_doublings+1,) + batch_shape.
    left_endpoints -= left_increments
    right_endpoints = left_endpoints + widths

    # Test if these end points lie outside of the slice.
    # Checks if the end points of the slice are outside the graph of the pdf.
    left_ep_values = tf.map_fn(target_log_prob, left_endpoints)
    right_ep_values = tf.map_fn(target_log_prob, right_endpoints)
    left_ok = left_ep_values < log_slice_heights
    right_ok = right_ep_values < log_slice_heights
    both_ok = left_ok & right_ok

    both_ok_f = tf.reshape(both_ok, [max_doublings + 1, -1])

    best_interval_idx = _find_best_interval_idx(
        tf.cast(both_ok_f, dtype=tf.int32))

    # Formats the above index as required to use with gather_nd.
    point_index_gather = tf.stack(
        [best_interval_idx,
         tf.range(tf.size(input=best_interval_idx))],
        axis=1,
        name='point_index_gather')
    left_ep_f = tf.reshape(left_endpoints, [max_doublings + 1, -1])
    right_ep_f = tf.reshape(right_endpoints, [max_doublings + 1, -1])
    # The x values of the uppper and lower bounds of the slices for each chain.
    lower_bounds = tf.reshape(tf.gather_nd(left_ep_f, point_index_gather),
                              batch_shape)
    upper_bounds = tf.reshape(tf.gather_nd(right_ep_f, point_index_gather),
                              batch_shape)
    both_ok = tf.reduce_any(input_tensor=both_ok, axis=0)
    return upper_bounds, lower_bounds, both_ok

def _test_acceptance(x_initial, target_log_prob, decided, log_slice_heights,
                     x_proposed, step_size, lower_bounds, upper_bounds,
                     name=None):
  """Ensures the chosen point does not violate reversibility.

    Implements Fig 6 of Neal 2003 page 717, which checks that the path from the
    existing point to the new point would also have been possible in reverse.
    This is done by checking that the algorithm would not have been terminated
    before reaching the old point.

  Args:
    x_initial: A tensor of any shape and real dtype. The initial positions of
      the chains. This function assumes that all the dimensions of `x_initial`
      are batch dimensions (i.e. the event shape is `[]`).
    target_log_prob: Callable accepting a tensor like `x_initial` and returning
      a tensor containing the log density at that point of the same shape.
    decided: A `tf.bool` tensor of the same shape as `x_initial`. Indicates
      whether the acceptance has already been decided. A point is tested only
      if `decided` for that point is False.
    log_slice_heights: Tensor of the same shape and dtype as the return value
      of `target_log_prob` when applied to `x_initial`. The log of the height of
      the chosen slice.
    x_proposed: A tensor of the same shape and dtype as `x_initial`. The
      proposed points.
    step_size: A tensor of shape and dtype compatible with `x_initial`. The min
      interval size in the doubling algorithm.
    lower_bounds: Tensor of same shape and dtype as `x_initial`. Slice lower
      bounds for each chain.
    upper_bounds: Tensor of same shape and dtype as `x_initial`. Slice upper
      bounds for each chain.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    acceptable: A boolean tensor of same shape as `x_initial` indicating whether
      the proposed points are acceptable for reversibility or not.
  """
  with tf.compat.v1.name_scope(name, 'test_acceptance', [
      x_initial, decided, log_slice_heights, x_proposed, step_size,
      lower_bounds, upper_bounds
  ]):
    d = tf.zeros_like(x_initial, dtype=tf.bool)
    # Keeps track of points for which the loop has "effectively terminated".
    # Termination is when either their interval width has shrunk to the minimum
    # value (step_size) or if the point has already been rejected.
    def cond(_, decided, *ignored_args):  # pylint: disable=unused-argument
      # Continue until all the points have been decided.
      return ~tf.reduce_all(input_tensor=decided)

    acceptable = tf.ones_like(x_initial, dtype=tf.bool)
    def body(acceptable, decided, left, right, d):
      """Checks reversibility as described on P717 of Neal 2003."""
      midpoint = (left + right) / 2
      divided = (((x_initial < midpoint) & (x_proposed >= midpoint)) |
                 ((x_proposed < midpoint) & (x_initial >= midpoint)))
      next_d = d | divided
      next_right = tf.where(x_proposed < midpoint, midpoint, right)
      next_left = tf.where(x_proposed >= midpoint, midpoint, left)
      left_test = (log_slice_heights >= target_log_prob(next_left))
      right_test = (log_slice_heights >= target_log_prob(next_right))
      unacceptable = next_d & left_test & right_test
      # Logic here: For points which have not already been decided,
      # and are unacceptable, set acceptable to False. For others, let them
      # be as they were.
      now_decided = ~decided & unacceptable
      next_acceptable = tf.where(now_decided, ~unacceptable, acceptable)
      # Decided if (a) was already decided, or
      # (b) the new width is less than 1.1 step_size, or
      # (c) was marked unacceptable.
      next_decided = (decided | (next_right - next_left <= 1.1 * step_size) |
                      now_decided)
      return (next_acceptable, next_decided, next_left, next_right, next_d)

    return tf.while_loop(
        cond=cond,
        body=body,
        loop_vars=(acceptable, decided, lower_bounds, upper_bounds, d))[0]

def _sample_with_shrinkage(x_initial, target_log_prob, log_slice_heights,
                           step_size, lower_bounds, upper_bounds, seed=None,
                           name=None):
  """Samples from the slice by applying shrinkage for rejected points.

  Implements the one dimensional slice sampling algorithm of Neal (2003), with a
  doubling algorithm (Neal 2003 P715 Fig. 4), which doubles the size of the
  interval at each iteration and shrinkage (Neal 2003 P716 Fig. 5), which
  reduces the width of the slice when a selected point is rejected, by setting
  the relevant bound that that value. Randomly sampled points are checked for
  two criteria: that they lie within the slice and that they pass the
  acceptability check (Neal 2003 P717 Fig. 6), which tests that the new state
  could have generated the previous one.

  Args:
    x_initial: A tensor of any shape. The initial positions of the chains. This
      function assumes that all the dimensions of `x_initial` are batch
      dimensions (i.e. the event shape is `[]`).
    target_log_prob: Callable accepting a tensor like `x_initial` and returning
      a tensor containing the log density at that point of the same shape.
    log_slice_heights: Tensor of the same shape and dtype as the return value
      of `target_log_prob` when applied to `x_initial`. The log of the height of
      the chosen slice.
    step_size: A tensor of shape and dtype compatible with `x_initial`. The min
      interval size in the doubling algorithm.
    lower_bounds: Tensor of same shape and dtype as `x_initial`. Slice lower
      bounds for each chain.
    upper_bounds: Tensor of same shape and dtype as `x_initial`. Slice upper
      bounds for each chain.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    x_proposed: A tensor of the same shape and dtype as `x_initial`. The next
      proposed state of the chain.
  """
  with tf.compat.v1.name_scope(
      name, 'sample_with_shrinkage',
      [x_initial, log_slice_heights, step_size, lower_bounds, upper_bounds]):
    seed_gen = distributions.SeedStream(seed, salt='_sample_with_shrinkage')
    # Keeps track of whether an acceptable sample has been found for the chain.
    found = tf.zeros_like(x_initial, dtype=tf.bool)
    cond = lambda found, *ignored_args: ~tf.reduce_all(input_tensor=found)
    x_next = tf.identity(x_initial)
    x_initial_shape = tf.shape(input=x_initial)
    x_initial_dtype = x_initial.dtype.base_dtype
    def _body(found, left, right, x_next):
      """Iterates until every chain has found a suitable next state."""
      proportions = tf.random.uniform(
          x_initial_shape, dtype=x_initial_dtype, seed=seed_gen())
      x_proposed = tf.where(~found, left + proportions * (right - left), x_next)
      accept_res = _test_acceptance(x_initial, target_log_prob=target_log_prob,
                                    decided=found,
                                    log_slice_heights=log_slice_heights,
                                    x_proposed=x_proposed, step_size=step_size,
                                    lower_bounds=left, upper_bounds=right)
      boundary_test = log_slice_heights < target_log_prob(x_proposed)
      can_accept = boundary_test & accept_res
      next_found = found | can_accept
      # Note that it might seem that we are moving the left and right end points
      # even if the point has been accepted (which is contrary to the stated
      # algorithm in Neal). However, this does not matter because the endpoints
      # for points that have been already accepted are not used again so it
      # doesn't matter what we do with them.
      next_left = tf.where(x_proposed < x_initial, x_proposed, left)
      next_right = tf.where(x_proposed >= x_initial, x_proposed, right)
      return next_found, next_left, next_right, x_proposed

    return tf.while_loop(
        cond=cond,
        body=_body,
        loop_vars=(found, lower_bounds, upper_bounds, x_next))[-1]

def slice_sampler_one_dim(target_log_prob, x_initial, step_size=0.01,
                          max_doublings=30, seed=None, name=None):
  """For a given x position in each Markov chain, returns the next x.

  Applies the one dimensional slice sampling algorithm as defined in Neal (2003)
  to an input tensor x of shape (num_chains,) where num_chains is the number of
  simulataneous Markov chains, and returns the next tensor x of shape
  (num_chains,) when these chains are evolved by the slice sampling algorithm.

  Args:
    target_log_prob: Callable accepting a tensor like `x_initial` and returning
      a tensor containing the log density at that point of the same shape.
    x_initial: A tensor of any shape. The initial positions of the chains. This
      function assumes that all the dimensions of `x_initial` are batch
      dimensions (i.e. the event shape is `[]`).
    step_size: A tensor of shape and dtype compatible with `x_initial`. The min
      interval size in the doubling algorithm.
    max_doublings: Scalar tensor of dtype `tf.int32`. The maximum number of
      doublings to try to find the slice bounds.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    retval: A tensor of the same shape and dtype as `x_initial`. The next state
      of the Markov chain.
    next_target_log_prob: The target log density evaluated at `retval`.
    bounds_satisfied: A tensor of bool dtype and shape batch dimensions.
    upper_bounds: Tensor of the same shape and dtype as `x_initial`. The upper
      bounds for the slice found.
    lower_bounds: Tensor of the same shape and dtype as `x_initial`. The lower
      bounds for the slice found.
  """
  with tf.compat.v1.name_scope(name, 'slice_sampler_one_dim',
                               [x_initial, step_size, max_doublings]):
    x_initial = tf.convert_to_tensor(value=x_initial)
    # Obtain the input dtype of the array.
    dtype = x_initial.dtype.base_dtype
    # Select the height of the slice. Tensor of shape x_initial.shape.
    log_slice_heights = target_log_prob(x_initial) - tf.random.gamma(
        tf.shape(input=x_initial), alpha=1, dtype=dtype, seed=seed)
    # Given the above x and slice heights, compute the bounds of the slice for
    # each chain.
    upper_bounds, lower_bounds, bounds_satisfied = slice_bounds_by_doubling(
        x_initial, target_log_prob, log_slice_heights, max_doublings, step_size,
        seed=seed)
    retval = _sample_with_shrinkage(x_initial, target_log_prob=target_log_prob,
                                    log_slice_heights=log_slice_heights,
                                    step_size=step_size,
                                    lower_bounds=lower_bounds,
                                    upper_bounds=upper_bounds, seed=seed)
    return (retval, target_log_prob(retval), bounds_satisfied,
            upper_bounds, lower_bounds)


# 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.
# ============================================================================
"""Markov chain Monte Carlo driver, `sample_chain_annealed_importance_chain`."""

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

import collections
# Dependency imports
import numpy as np

import tensorflow as tf
from tensorflow_probability.python.mcmc.internal import util as mcmc_util

__all__ = [
    "sample_annealed_importance_chain",
]

AISResults = collections.namedtuple(
    "AISResults",
    [
        "proposal_log_prob",
        "target_log_prob",
        "inner_results",
    ])

def sample_annealed_importance_chain(
    num_steps,
    proposal_log_prob_fn,
    target_log_prob_fn,
    current_state,
    make_kernel_fn,
    parallel_iterations=10,
    name=None):
  """Runs annealed importance sampling (AIS) to estimate normalizing constants.

  This function uses an MCMC transition operator (e.g., Hamiltonian Monte Carlo)
  to sample from a series of distributions that slowly interpolates between
  an initial "proposal" distribution:

  `exp(proposal_log_prob_fn(x) - proposal_log_normalizer)`

  and the target distribution:

  `exp(target_log_prob_fn(x) - target_log_normalizer)`,

  accumulating importance weights along the way. The product of these
  importance weights gives an unbiased estimate of the ratio of the
  normalizing constants of the initial distribution and the target
  distribution:

  `E[exp(ais_weights)] = exp(target_log_normalizer - proposal_log_normalizer)`.

  Note: When running in graph mode, `proposal_log_prob_fn` and
  `target_log_prob_fn` are called exactly three times (although this may be
  reduced to two times in the future).

  Args:
    num_steps: Integer number of Markov chain updates to run. More
      iterations means more expense, but smoother annealing between q
      and p, which in turn means exponentially lower variance for the
      normalizing constant estimator.
    proposal_log_prob_fn: Python callable that returns the log density of the
      initial distribution.
    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.
    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))`.
    make_kernel_fn: Python `callable` which returns a `TransitionKernel`-like
      object. Must take one argument representing the `TransitionKernel`'s
      `target_log_prob_fn`. The `target_log_prob_fn` argument represents the
      `TransitionKernel`'s target log distribution.  Note:
      `sample_annealed_importance_chain` creates a new `target_log_prob_fn`
      which is an interpolation between the supplied `target_log_prob_fn` and
      `proposal_log_prob_fn`; it is this interpolated function which is used as
      an argument to `make_kernel_fn`.
    parallel_iterations: The number of iterations allowed to run in parallel.
        It must be a positive integer. See `tf.while_loop` for more details.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "sample_annealed_importance_chain").

  Returns:
    next_state: `Tensor` or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at the final iteration. Has same shape as
      input `current_state`.
    ais_weights: Tensor with the estimated weight(s). Has shape matching
      `target_log_prob_fn(current_state)`.
    kernel_results: `collections.namedtuple` of internal calculations used to
      advance the chain.

  #### Examples

  ##### Estimate the normalizing constant of a log-gamma distribution.

  ```python
  tfd = tfp.distributions

  # Run 100 AIS chains in parallel
  num_chains = 100
  dims = 20
  dtype = np.float32

  proposal = tfd.MultivatiateNormalDiag(
     loc=tf.zeros([dims], dtype=dtype))

  target = tfd.TransformedDistribution(
    distribution=tfd.Gamma(concentration=dtype(2),
                           rate=dtype(3)),
    bijector=tfp.bijectors.Invert(tfp.bijectors.Exp()),
    event_shape=[dims])

  chains_state, ais_weights, kernels_results = (
      tfp.mcmc.sample_annealed_importance_chain(
          num_steps=1000,
          proposal_log_prob_fn=proposal.log_prob,
          target_log_prob_fn=target.log_prob,
          current_state=proposal.sample(num_chains),
          make_kernel_fn=lambda tlp_fn: tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=tlp_fn,
            step_size=0.2,
            num_leapfrog_steps=2)))

  log_estimated_normalizer = (tf.reduce_logsumexp(ais_weights)
                              - np.log(num_chains))
  log_true_normalizer = tf.lgamma(2.) - 2. * tf.log(3.)
  ```

  ##### Estimate marginal likelihood of a Bayesian regression model.

  ```python
  tfd = tfp.distributions

  def make_prior(dims, dtype):
    return tfd.MultivariateNormalDiag(
        loc=tf.zeros(dims, dtype))

  def make_likelihood(weights, x):
    return tfd.MultivariateNormalDiag(
        loc=tf.tensordot(weights, x, axes=[[0], [-1]]))

  # Run 100 AIS chains in parallel
  num_chains = 100
  dims = 10
  dtype = np.float32

  # Make training data.
  x = np.random.randn(num_chains, dims).astype(dtype)
  true_weights = np.random.randn(dims).astype(dtype)
  y = np.dot(x, true_weights) + np.random.randn(num_chains)

  # Setup model.
  prior = make_prior(dims, dtype)
  def target_log_prob_fn(weights):
    return prior.log_prob(weights) + make_likelihood(weights, x).log_prob(y)

  proposal = tfd.MultivariateNormalDiag(
      loc=tf.zeros(dims, dtype))

  weight_samples, ais_weights, kernel_results = (
      tfp.mcmc.sample_annealed_importance_chain(
        num_steps=1000,
        proposal_log_prob_fn=proposal.log_prob,
        target_log_prob_fn=target_log_prob_fn
        current_state=tf.zeros([num_chains, dims], dtype),
        make_kernel_fn=lambda tlp_fn: tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=tlp_fn,
          step_size=0.1,
          num_leapfrog_steps=2)))
  log_normalizer_estimate = (tf.reduce_logsumexp(ais_weights)
                             - np.log(num_chains))
  ```

  """
  with tf.compat.v1.name_scope(name, "sample_annealed_importance_chain",
                               [num_steps, current_state]):
    num_steps = tf.convert_to_tensor(
        value=num_steps, dtype=tf.int32, name="num_steps")
    if mcmc_util.is_list_like(current_state):
      current_state = [
          tf.convert_to_tensor(value=s, name="current_state")
          for s in current_state
      ]
    else:
      current_state = tf.convert_to_tensor(
          value=current_state, name="current_state")

    def _make_convex_combined_log_prob_fn(iter_):
      def _fn(*args):
        p = tf.identity(proposal_log_prob_fn(*args), name="proposal_log_prob")
        t = tf.identity(target_log_prob_fn(*args), name="target_log_prob")
        dtype = p.dtype.base_dtype
        beta = tf.cast(iter_ + 1, dtype) / tf.cast(num_steps, dtype)
        return tf.identity(beta * t + (1. - beta) * p,
                           name="convex_combined_log_prob")
      return _fn

    def _loop_body(iter_, ais_weights, current_state, kernel_results):
      """Closure which implements `tf.while_loop` body."""
      x = (current_state if mcmc_util.is_list_like(current_state)
           else [current_state])
      proposal_log_prob = proposal_log_prob_fn(*x)
      target_log_prob = target_log_prob_fn(*x)
      ais_weights += ((target_log_prob - proposal_log_prob) /
                      tf.cast(num_steps, ais_weights.dtype))
      kernel = make_kernel_fn(_make_convex_combined_log_prob_fn(iter_))
      next_state, inner_results = kernel.one_step(
          current_state, kernel_results.inner_results)
      kernel_results = AISResults(
          proposal_log_prob=proposal_log_prob,
          target_log_prob=target_log_prob,
          inner_results=inner_results,
      )
      return [iter_ + 1, ais_weights, next_state, kernel_results]

    def _bootstrap_results(init_state):
      """Creates first version of `previous_kernel_results`."""
      kernel = make_kernel_fn(_make_convex_combined_log_prob_fn(iter_=0))
      inner_results = kernel.bootstrap_results(init_state)

      convex_combined_log_prob = inner_results.accepted_results.target_log_prob
      dtype = convex_combined_log_prob.dtype.as_numpy_dtype
      shape = tf.shape(input=convex_combined_log_prob)
      proposal_log_prob = tf.fill(shape, dtype(np.nan),
                                  name="bootstrap_proposal_log_prob")
      target_log_prob = tf.fill(shape, dtype(np.nan),
                                name="target_target_log_prob")

      return AISResults(
          proposal_log_prob=proposal_log_prob,
          target_log_prob=target_log_prob,
          inner_results=inner_results,
      )

    previous_kernel_results = _bootstrap_results(current_state)
    inner_results = previous_kernel_results.inner_results

    ais_weights = tf.zeros(
        shape=tf.broadcast_dynamic_shape(
            tf.shape(input=inner_results.proposed_results.target_log_prob),
            tf.shape(input=inner_results.accepted_results.target_log_prob)),
        dtype=inner_results.proposed_results.target_log_prob.dtype.base_dtype)

    [_, ais_weights, current_state, kernel_results] = tf.while_loop(
        cond=lambda iter_, *args: iter_ < num_steps,
        body=_loop_body,
        loop_vars=[
            np.int32(0),  # iter_
            ais_weights,
            current_state,
            previous_kernel_results,
        ],
        parallel_iterations=parallel_iterations)

    return [current_state, ais_weights, kernel_results]


# 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.
# ============================================================================
"""Transformations of Edward2 programs."""

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

import inspect
import six
import tensorflow as tf

from tensorflow_probability.python.edward2.interceptor import interceptable
from tensorflow_probability.python.edward2.interceptor import interception

__all__ = [
    "make_log_joint_fn",
    "make_value_setter"
]

def make_value_setter(**model_kwargs):
  """Creates a value-setting interceptor.

  This function creates an interceptor that sets values of Edward2 random
  variable objects. This is useful for a range of tasks, including conditioning
  on observed data, sampling from posterior predictive distributions, and as a
  building block of inference primitives such as computing log joint
  probabilities (see examples below).

  Args:
    **model_kwargs: dict of str to Tensor. Keys are the names of random
      variables in the model to which this interceptor is being applied. Values
      are Tensors to set their value to. Variables not included in this dict
      will not be set and will maintain their existing value semantics (by
      default, a sample from the parent-conditional distribution).

  Returns:
    set_values: function that sets the value of intercepted ops.

  #### Examples

  Consider for illustration a model with latent `z` and
  observed `x`, and a corresponding trainable posterior model:

  ```python
  num_observations = 10
  def model():
    z = ed.Normal(loc=0, scale=1., name='z')  # log rate
    x = ed.Poisson(rate=tf.exp(z) * tf.ones(num_observations), name='x')
    return x

  def variational_model():
    return ed.Normal(loc=tf.Variable(0.),
                     scale=tf.nn.softplus(tf.Variable(-4.)),
                     name='z')  # for simplicity, match name of the model RV.
  ```

  We can use a value-setting interceptor to condition the model on observed
  data. This approach is slightly more cumbersome than that of partially
  evaluating the complete log-joint function, but has the potential advantage
  that it returns a new model callable, which may be used to sample downstream
  variables, passed into additional transformations, etc.

  ```python
  x_observed = np.array([6, 3, 1, 8, 7, 0, 6, 4, 7, 5])
  def observed_model():
    with ed.interception(make_value_setter(x=x_observed)):
      model()
  observed_log_joint_fn = ed.make_log_joint_fn(observed_model)

  # After fixing 'x', the observed log joint is now only a function of 'z'.
  # This enables us to define a variational lower bound,
  # `E_q[ log p(x, z) - log q(z)]`, simply by evaluating the observed and
  # variational log joints at variational samples.
  variational_log_joint_fn = ed.make_log_joint_fn(variational_model)
  with ed.tape() as variational_sample:  # Sample trace from variational model.
    variational_model()
  elbo_loss = -(observed_log_joint_fn(**variational_sample) -
                variational_log_joint_fn(**variational_sample))
  ```

  After performing inference by minimizing the variational loss, a value-setting
  interceptor enables simulation from the posterior predictive distribution:

  ```python
  with ed.tape() as posterior_samples:  # tape is a map {rv.name : rv}
    variational_model()
  with ed.interception(ed.make_value_setter(**posterior_samples)):
    x = model()
  # x is a sample from p(X | Z = z') where z' ~ q(z) (the variational model)
  ```

  As another example, using a value setter inside of `ed.tape` enables
  computing the log joint probability, by setting all variables to
  posterior values and then accumulating the log probs of those values under
  the induced parent-conditional distributions. This is one way that we could
  have implemented `ed.make_log_joint_fn`:

  ```python
  def make_log_joint_fn_demo(model):
    def log_joint_fn(**model_kwargs):
      with ed.tape() as model_tape:
        with ed.make_value_setter(**model_kwargs):
          model()

      # accumulate sum_i log p(X_i = x_i | X_{:i-1} = x_{:i-1})
      log_prob = 0.
      for rv in model_tape.values():
        log_prob += tf.reduce_sum(rv.log_prob(rv.value))

      return log_prob
    return log_joint_fn
  ```

  """
  def set_values(f, *args, **kwargs):
    """Sets random variable values to its aligned value."""
    name = kwargs.get("name")
    if name in model_kwargs:
      kwargs["value"] = model_kwargs[name]
    return interceptable(f)(*args, **kwargs)
  return set_values

def make_log_joint_fn(model):
  """Takes Edward probabilistic program and returns its log joint function.

  Args:
    model: Python callable which executes the generative process of a
      computable probability distribution using `ed.RandomVariable`s.

  Returns:
    A log-joint probability function. Its inputs are `model`'s original inputs
    and random variables which appear during the program execution. Its output
    is a scalar tf.Tensor.

  #### Examples

  Below we define Bayesian logistic regression as an Edward program,
  representing the model's generative process. We apply `make_log_joint_fn` in
  order to represent the model in terms of its joint probability function.

  ```python
  from tensorflow_probability import edward2 as ed

  def logistic_regression(features):
    coeffs = ed.Normal(loc=0., scale=1.,
                       sample_shape=features.shape[1], name="coeffs")
    outcomes = ed.Bernoulli(logits=tf.tensordot(features, coeffs, [[1], [0]]),
                            name="outcomes")
    return outcomes

  log_joint = ed.make_log_joint_fn(logistic_regression)

  features = tf.random_normal([3, 2])
  coeffs_value = tf.random_normal([2])
  outcomes_value = tf.round(tf.random_uniform([3]))
  output = log_joint(features, coeffs=coeffs_value, outcomes=outcomes_value)
  ```

  """
  def log_joint_fn(*args, **kwargs):
    """Log-probability of inputs according to a joint probability distribution.

    Args:
      *args: Positional arguments. They are the model's original inputs and can
        alternatively be specified as part of `kwargs`.
      **kwargs: Keyword arguments, where for each key-value pair `k` and `v`,
        `v` is passed as a `value` to the random variable(s) whose keyword
        argument `name` during construction is equal to `k`.

    Returns:
      Scalar tf.Tensor, which represents the model's log-probability summed
      over all Edward random variables and their dimensions.

    Raises:
      TypeError: If a random variable in the model has no specified value in
        `**kwargs`.
    """
    log_probs = []

    def interceptor(rv_constructor, *rv_args, **rv_kwargs):
      """Overrides a random variable's `value` and accumulates its log-prob."""
      # Set value to keyword argument indexed by `name` (an input tensor).
      rv_name = rv_kwargs.get("name")
      if rv_name is None:
        raise KeyError("Random variable constructor {} has no name "
                       "in its arguments.".format(rv_constructor.__name__))

      # If no value is explicitly passed in for an RV, default to the value
      # from the RV constructor. This may have been set explicitly by the user
      # or forwarded from a lower-level interceptor.
      previously_specified_value = rv_kwargs.get("value")
      value = kwargs.get(rv_name, previously_specified_value)
      if value is None:
        raise LookupError("Keyword argument specifying value for {} is "
                          "missing.".format(rv_name))
      rv_kwargs["value"] = value

      rv = rv_constructor(*rv_args, **rv_kwargs)
      log_prob = tf.reduce_sum(input_tensor=rv.distribution.log_prob(rv.value))
      log_probs.append(log_prob)
      return rv

    model_kwargs = _get_function_inputs(model, kwargs)
    with interception(interceptor):
      model(*args, **model_kwargs)
    log_prob = sum(log_probs)
    return log_prob
  return log_joint_fn

def _get_function_inputs(f, src_kwargs):
  """Filters inputs to be compatible with function `f`'s signature.

  Args:
    f: Function according to whose input signature we filter arguments.
    src_kwargs: Keyword arguments to filter according to `f`.

  Returns:
    kwargs: Dict of key-value pairs in `src_kwargs` which exist in `f`'s
      signature.
  """
  if hasattr(f, "_func"):  # functions returned by tf.make_template
    f = f._func  # pylint: disable=protected-access

  try:  # getargspec was deprecated in Python 3.6
    argspec = inspect.getfullargspec(f)
  except AttributeError:
    argspec = inspect.getargspec(f)

  fkwargs = {k: v for k, v in six.iteritems(src_kwargs) if k in argspec.args}
  return fkwargs