# 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.
# ============================================================================
"""Internal utility functions for implementing TransitionKernels."""

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

import functools
import warnings

# Dependency imports
import numpy as np
import tensorflow as tf

from tensorflow_probability.python.math.gradient import value_and_gradient as tfp_math_value_and_gradients

from tensorflow.python.ops import control_flow_util  # pylint: disable=g-direct-tensorflow-import
from tensorflow.python.util import deprecation  # pylint: disable=g-direct-tensorflow-import

__all__ = [
    'choose',
    'is_list_like',
    'is_namedtuple_like',
    'make_innermost_setter',
    'make_innermost_getter',
    'make_name',
    'maybe_call_fn_and_grads',
    'safe_sum',
    'set_doc',
    'smart_for_loop',
    'trace_scan',
    'enable_store_parameters_in_results',
    'warn_if_parameters_are_not_simple_tensors',
]

def is_list_like(x):
  """Helper which returns `True` if input is `list`-like."""
  return isinstance(x, (tuple, list))

def is_namedtuple_like(x):
  """Helper which returns `True` if input is `collections.namedtuple`-like."""
  try:
    for fn in x._fields:
      _ = getattr(x, fn)
    return True
  except AttributeError:
    return False

def make_name(super_name, default_super_name, sub_name):
  """Helper which makes a `str` name; useful for tf.compat.v1.name_scope."""
  name = super_name if super_name is not None else default_super_name
  if sub_name is not None:
    name += '_' + sub_name
  return name

def _choose_base_case(is_accepted,
                      accepted,
                      rejected,
                      name=None):
  """Helper to `choose` which expand_dims `is_accepted` and applies tf.where."""
  def _expand_is_accepted_like(x):
    """Helper to expand `is_accepted` like the shape of some input arg."""
    with tf.compat.v1.name_scope('expand_is_accepted_like'):
      expand_shape = tf.concat([
          tf.shape(input=is_accepted),
          tf.ones([tf.rank(x) - tf.rank(is_accepted)], dtype=tf.int32),
      ],
                               axis=0)
      multiples = tf.concat([
          tf.ones([tf.rank(is_accepted)], dtype=tf.int32),
          tf.shape(input=x)[tf.rank(is_accepted):],
      ],
                            axis=0)
      m = tf.tile(tf.reshape(is_accepted, expand_shape),
                  multiples)
      m.set_shape(m.shape.merge_with(x.shape))
      return m
  def _where(accepted, rejected):
    if accepted is rejected:
      return accepted
    accepted = tf.convert_to_tensor(value=accepted, name='accepted')
    rejected = tf.convert_to_tensor(value=rejected, name='rejected')
    r = tf.where(_expand_is_accepted_like(accepted), accepted, rejected)
    r.set_shape(r.shape.merge_with(accepted.shape.merge_with(rejected.shape)))
    return r

  with tf.compat.v1.name_scope(
      name, 'choose', values=[is_accepted, accepted, rejected]):
    if not is_list_like(accepted):
      return _where(accepted, rejected)
    return [(choose(is_accepted, a, r, name=name) if is_namedtuple_like(a)
             else _where(a, r))
            for a, r in zip(accepted, rejected)]

def choose(is_accepted, accepted, rejected, name=None):
  """Helper which expand_dims `is_accepted` then applies tf.where."""
  if not is_namedtuple_like(accepted):
    return _choose_base_case(is_accepted, accepted, rejected, name=name)
  if not isinstance(accepted, type(rejected)):
    raise TypeError('Type of `accepted` ({}) must be identical to '
                    'type of `rejected` ({})'.format(
                        type(accepted).__name__,
                        type(rejected).__name__))
  return type(accepted)(**dict(
      [(fn,
        choose(is_accepted,
               getattr(accepted, fn),
               getattr(rejected, fn),
               name=name))
       for fn in accepted._fields]))

def safe_sum(x, alt_value=-np.inf, name=None):
  """Elementwise adds list members, replacing non-finite results with alt_value.

  Typically the `alt_value` is chosen so the `MetropolisHastings`
  `TransitionKernel` always rejects the proposal.

  Args:
    x: Python `list` of `Tensors` to elementwise add.
    alt_value: Python scalar used to replace any elementwise sums which would
      otherwise be non-finite.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "safe_sum").

  Returns:
    safe_sum: `Tensor` representing the elementwise sum of list of `Tensor`s
      `x` or `alt_value` where sums are non-finite.

  Raises:
    TypeError: if `x` is not list-like.
    ValueError: if `x` is empty.
  """
  with tf.compat.v1.name_scope(name, 'safe_sum', [x, alt_value]):
    if not is_list_like(x):
      raise TypeError('Expected list input.')
    if not x:
      raise ValueError('Input should not be empty.')
    in_shape = x[0].shape
    x = tf.stack(x, axis=-1)
    x = tf.reduce_sum(input_tensor=x, axis=-1)
    alt_value = np.array(alt_value, x.dtype.as_numpy_dtype)
    alt_fill = tf.fill(tf.shape(input=x), value=alt_value)
    x = tf.where(tf.math.is_finite(x), x, alt_fill)
    x.set_shape(x.shape.merge_with(in_shape))
    return x

def set_doc(value):
  """Decorator to programmatically set a function docstring."""
  def _doc(func):
    func.__doc__ = value
    return func
  return _doc

def _value_and_gradients(fn, fn_arg_list, result=None, grads=None, name=None):
  """Helper to `maybe_call_fn_and_grads`."""
  with tf.compat.v1.name_scope(name, 'value_and_gradients',
                               [fn_arg_list, result, grads]):

    def _convert_to_tensor(x, name):
      ctt = lambda x_: x_ if x_ is None else tf.convert_to_tensor(
          value=x_, name=name)
      return [ctt(x_) for x_ in x] if is_list_like(x) else ctt(x)

    fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)
                   else [fn_arg_list])
    fn_arg_list = _convert_to_tensor(fn_arg_list, 'fn_arg')

    if result is None:
      result = fn(*fn_arg_list)
      if grads is None and tf.executing_eagerly():
        # Ensure we disable bijector cacheing in eager mode.
        # TODO(b/72831017): Remove this once bijector cacheing is fixed for
        # eager mode.
        fn_arg_list = [0 + x for x in fn_arg_list]

    result = _convert_to_tensor(result, 'fn_result')

    if grads is not None:
      grads = _convert_to_tensor(grads, 'fn_grad')
      return result, grads

    if is_list_like(result) and len(result) == len(fn_arg_list):
      # Compute the block diagonal of Jacobian.
      # TODO(b/79158574): Guard this calculation by an arg which explicitly
      # requests block diagonal Jacobian calculation.
      def fn_slice(i):
        """Needed to prevent `cell-var-from-loop` pylint warning."""
        return lambda x: fn(*(fn_arg_list[:i] + [x] + fn_arg_list[i+1:]))
      grads = [
          tfp_math_value_and_gradients(fn_slice(i), fn_arg_list[i])[1]
          for i in range(len(result))
      ]
    else:
      _, grads = tfp_math_value_and_gradients(fn, fn_arg_list)

    return result, grads

def maybe_call_fn_and_grads(fn,
                            fn_arg_list,
                            result=None,
                            grads=None,
                            check_non_none_grads=True,
                            name=None):
  """Calls `fn` and computes the gradient of the result wrt `args_list`."""
  with tf.compat.v1.name_scope(name, 'maybe_call_fn_and_grads',
                               [fn_arg_list, result, grads]):
    fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)
                   else [fn_arg_list])
    result, grads = _value_and_gradients(fn, fn_arg_list, result, grads)
    if not all(r.dtype.is_floating
               for r in (result if is_list_like(result) else [result])):  # pylint: disable=superfluous-parens
      raise TypeError('Function result must be a `Tensor` with `float` '
                      '`dtype`.')
    if len(fn_arg_list) != len(grads):
      raise ValueError('Function args must be in one-to-one correspondence '
                       'with grads.')
    if check_non_none_grads and any(g is None for g in grads):
      raise ValueError('Encountered `None` gradient.\n'
                       '  fn_arg_list: {}\n'
                       '  grads: {}'.format(fn_arg_list, grads))
    return result, grads

def smart_for_loop(loop_num_iter, body_fn, initial_loop_vars,
                   parallel_iterations=10, name=None):
  """Construct a for loop, preferring a python loop if `n` is staticaly known.

  Given `loop_num_iter` and `body_fn`, return an op corresponding to executing
  `body_fn` `loop_num_iter` times, feeding previous outputs of `body_fn` into
  the next iteration.

  If `loop_num_iter` is statically known, the op is constructed via python for
  loop, and otherwise a `tf.while_loop` is used.

  Args:
    loop_num_iter: `Integer` `Tensor` representing the number of loop
      iterations.
    body_fn: Callable to be executed `loop_num_iter` times.
    initial_loop_vars: Listlike object of `Tensors` to be passed in to
      `body_fn`'s first execution.
    parallel_iterations: The number of iterations allowed to run in parallel.
      It must be a positive integer. See `tf.while_loop` for more details.
      Default value: `10`.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "smart_for_loop").
  Returns:
    result: `Tensor` representing applying `body_fn` iteratively `n` times.
  """
  with tf.compat.v1.name_scope(name, 'smart_for_loop',
                               [loop_num_iter, initial_loop_vars]):
    loop_num_iter_ = tf.get_static_value(loop_num_iter)
    if (loop_num_iter_ is None or tf.executing_eagerly() or
        control_flow_util.GraphOrParentsInXlaContext(
            tf.compat.v1.get_default_graph())):
      # Cast to int32 to run the comparison against i in host memory,
      # where while/LoopCond needs it.
      loop_num_iter = tf.cast(loop_num_iter, dtype=tf.int32)
      return tf.while_loop(
          cond=lambda i, *args: i < loop_num_iter,
          body=lambda i, *args: [i + 1] + list(body_fn(*args)),
          loop_vars=[np.int32(0)] + initial_loop_vars,
          parallel_iterations=parallel_iterations
      )[1:]
    result = initial_loop_vars
    for _ in range(loop_num_iter_):
      result = body_fn(*result)
    return result

def trace_scan(loop_fn,
               initial_state,
               elems,
               trace_fn,
               parallel_iterations=10,
               name=None):
  """A simplified version of `tf.scan` that has configurable tracing.

  This function repeatedly calls `loop_fn(state, elem)`, where `state` is the
  `initial_state` during the first iteration, and the return value of `loop_fn`
  for every iteration thereafter. `elem` is a slice of `elements` along the
  first dimension, accessed in order. Additionally, it calls `trace_fn` on the
  return value of `loop_fn`. The `Tensor`s in return values of `trace_fn` are
  stacked and returned from this function, such that the first dimension of
  those `Tensor`s matches the size of `elems`.

  Args:
    loop_fn: A callable that takes in a `Tensor` or a nested collection of
      `Tensor`s with the same structure as `initial_state`, a slice of `elems`
      and returns the same structure as `initial_state`.
    initial_state: A `Tensor` or a nested collection of `Tensor`s passed to
      `loop_fn` in the first iteration.
    elems: A `Tensor` that is split along the first dimension and each element
      of which is passed to `loop_fn`.
    trace_fn: A callable that takes in the return value of `loop_fn` and returns
      a `Tensor` or a nested collection of `Tensor`s.
    parallel_iterations: Passed to the internal `tf.while_loop`.
    name: Name scope used in this function. Default: 'trace_scan'.

  Returns:
    final_state: The final return value of `loop_fn`.
    trace: The same structure as the return value of `trace_fn`, but with each
      `Tensor` being a stack of the corresponding `Tensors` in the return value
      of `trace_fn` for each slice of `elems`.
  """
  with tf.compat.v1.name_scope(
      name, 'trace_scan', [initial_state, elems]), tf.compat.v1.variable_scope(
          tf.compat.v1.get_variable_scope()) as vs:
    if vs.caching_device is None and not tf.executing_eagerly():
      vs.set_caching_device(lambda op: op.device)

    initial_state = tf.nest.map_structure(
        lambda x: tf.convert_to_tensor(value=x, name='initial_state'),
        initial_state)
    elems = tf.convert_to_tensor(value=elems, name='elems')

    static_length = elems.shape[0]
    if tf.compat.dimension_value(static_length) is None:
      length = tf.shape(input=elems)[0]
    else:
      length = tf.convert_to_tensor(
          value=static_length, dtype=tf.int32, name='length')

    # This is an TensorArray in part because of XLA, which had trouble with
    # non-statically known indices. I.e. elems[i] errored, but
    # elems_array.read(i) worked.
    elems_array = tf.TensorArray(
        elems.dtype, size=length, element_shape=elems.shape[1:])
    elems_array = elems_array.unstack(elems)

    trace_arrays = tf.nest.map_structure(
        lambda x: tf.TensorArray(x.dtype, size=length, element_shape=x.shape),
        trace_fn(initial_state))

    def _body(i, state, trace_arrays):
      state = loop_fn(state, elems_array.read(i))
      trace_arrays = tf.nest.pack_sequence_as(trace_arrays, [
          a.write(i, v) for a, v in zip(
              tf.nest.flatten(trace_arrays), tf.nest.flatten(trace_fn(state)))
      ])
      return i + 1, state, trace_arrays

    _, final_state, trace_arrays = tf.while_loop(
        cond=lambda i, *args: i < length,
        body=_body,
        loop_vars=(0, initial_state, trace_arrays),
        parallel_iterations=parallel_iterations)

    stacked_trace = tf.nest.map_structure(lambda x: x.stack(), trace_arrays)

    # Restore the static length if we know it.
    def _merge_static_length(x):
      x.set_shape(tf.TensorShape(static_length).concatenate(x.shape[1:]))
      return x

    stacked_trace = tf.nest.map_structure(_merge_static_length, stacked_trace)
    return final_state, stacked_trace

def make_innermost_setter(setter):
  """Wraps a setter so it applies to the inner-most results in `kernel_results`.

  The wrapped setter unwraps `kernel_results` and applies `setter` to the first
  results without an `inner_results` attribute.

  Args:
    setter: A callable that takes the kernel results as well as some `*args` and
      `**kwargs` and returns a modified copy of those kernel results.

  Returns:
    new_setter: A wrapped `setter`.
  """

  @functools.wraps(setter)
  def _new_setter(kernel_results, *args, **kwargs):
    """Wrapped setter."""
    results_stack = []
    while hasattr(kernel_results, 'inner_results'):
      results_stack.append(kernel_results)
      kernel_results = kernel_results.inner_results

    new_kernel_results = setter(kernel_results, *args, **kwargs)
    for outer_results in reversed(results_stack):
      new_kernel_results = outer_results._replace(
          inner_results=new_kernel_results)

    return new_kernel_results

  return _new_setter

def make_innermost_getter(getter):
  """Wraps a getter so it applies to the inner-most results in `kernel_results`.

  The wrapped getter unwraps `kernel_results` and returns the return value of
  `getter` called with the first results without an `inner_results` attribute.

  Args:
    getter: A callable that takes Kernel results and returns some value.

  Returns:
    new_getter: A wrapped `getter`.
  """

  @functools.wraps(getter)
  def _new_getter(kernel_results, *args, **kwargs):
    """Wrapped getter."""
    results_stack = []
    while hasattr(kernel_results, 'inner_results'):
      results_stack.append(kernel_results)
      kernel_results = kernel_results.inner_results

    return getter(kernel_results, *args, **kwargs)

  return _new_getter

def enable_store_parameters_in_results(kernel):
  """Enables the `store_parameters_in_results` parameter in a chain of kernels.

  This is a temporary utility for use during the transition period of the
  parameter storage methods.

  Args:
    kernel: A TransitionKernel.

  Returns:
    kernel: The same kernel, but recreated with `store_parameters_in_results`
        recursively set to `True` in its parameters and its inner kernels (as
        appropriate).
  """
  kernel_stack = []
  while hasattr(kernel, 'parameters') and 'inner_kernel' in kernel.parameters:
    kernel_stack.append(kernel)
    kernel = kernel.parameters['inner_kernel']

  def _recreate_kernel(kernel, parameters):
    new_parameters = kernel.parameters.copy()
    new_parameters.update(parameters)
    if 'store_parameters_in_results' in new_parameters:
      new_parameters['store_parameters_in_results'] = True
    with deprecation.silence():
      return type(kernel)(**new_parameters)

  if hasattr(kernel, 'parameters'):
    kernel = _recreate_kernel(kernel, {})

  for outer_kernel in reversed(kernel_stack):
    outer_kernel = _recreate_kernel(outer_kernel, {'inner_kernel': kernel})
    kernel = outer_kernel

  return kernel

def _is_tensor_like(param):
  if is_list_like(param):
    return all([_is_tensor_like(p) for p in param])
  return isinstance(param, tf.Tensor) or (np.array(param).dtype != np.object)

def warn_if_parameters_are_not_simple_tensors(params_dict):
  for param_name, param in params_dict.items():
    if not _is_tensor_like(param):
      warnings.warn(
          '`{}` is not a `tf.Tensor`, Python number, or Numpy array. If this '
          'parameter is mutable (e.g., a `tf.Variable`), then the '
          'behavior implied by `store_parameters_in_results` will silently '
          'change on 2019-08-01. Please consult the docstring for '
          '`store_parameters_in_results` details and use '
          '`store_parameters_in_results=True` to silence this warning.'.format(
              param_name))


# 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.
# ============================================================================
"""Reshape bijectors."""

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

# Dependency imports
import numpy as np

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.bijectors import bijector
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    'Reshape',
]

class Reshape(bijector.Bijector):
  """Reshapes the `event_shape` of a `Tensor`.

  The semantics generally follow that of `tf.reshape()`, with
  a few differences:

  * The user must provide both the input and output shape, so that
    the transformation can be inverted. If an input shape is not
    specified, the default assumes a vector-shaped input, i.e.,
    event_shape_in = (-1,).
  * The `Reshape` bijector automatically broadcasts over the leftmost
    dimensions of its input (`sample_shape` and `batch_shape`); only
    the rightmost `event_ndims_in` dimensions are reshaped. The
    number of dimensions to reshape is inferred from the provided
    `event_shape_in` (`event_ndims_in = len(event_shape_in)`).

  Example usage:

  ```python
  r = tfp.bijectors.Reshape(event_shape_out=[1, -1])

  r.forward([3., 4.])    # shape [2]
  # ==> [[3., 4.]]       # shape [1, 2]

  r.forward([[1., 2.], [3., 4.]])  # shape [2, 2]
  # ==> [[[1., 2.]],
  #      [[3., 4.]]]   # shape [2, 1, 2]

  r.inverse([[3., 4.]])  # shape [1,2]
  # ==> [3., 4.]         # shape [2]

  r.forward_log_det_jacobian(any_value)
  # ==> 0.

  r.inverse_log_det_jacobian(any_value)
  # ==> 0.
  ```

  Note: we had to make a tricky-to-describe policy decision, which we attempt to
  summarize here. At instantiation time and class method invocation time, we
  validate consistency of class-level and method-level arguments. Note that
  since the class-level arguments may be unspecified until graph execution time,
  we had the option of deciding between two validation policies. One was,
  roughly, "the earliest, most statically specified arguments take precedence".
  The other was "method-level arguments must be consistent with class-level
  arguments". The former policy is in a sense more optimistic about user intent,
  and would enable us, in at least one particular case [1], to perform
  additional inference about resulting shapes. We chose the latter policy, as it
  is simpler to implement and a bit easier to articulate.

  [1] The case in question is exemplified in the following snippet:

  ```python
  bijector = tfp.bijectors.Reshape(
    event_shape_out=tf.placeholder(dtype=tf.int32, shape=[1]),
    event_shape_in= tf.placeholder(dtype=tf.int32, shape=[3]),
    validate_args=True)

  bijector.forward_event_shape(tf.TensorShape([5, 2, 3, 7]))
  # Chosen policy    ==> (5, None)
  # Alternate policy ==> (5, 42)
  ```

  In the chosen policy, since we don't know what `event_shape_in/out` are at the
  time of the call to `forward_event_shape`, we simply fill in everything we
  *do* know, which is that the last three dims will be replaced with
  "something".

  In the alternate policy, we would assume that the intention must be to reshape
  `[5, 2, 3, 7]` such that the last three dims collapse to one, which is only
  possible if the resulting shape is `[5, 42]`.

  Note that the above is the *only* case in which we could do such inference; if
  the output shape has more than 1 dim, we can't infer anything. E.g., we would
  have

  ```python
  bijector = tfp.bijectors.Reshape(
    event_shape_out=tf.placeholder(dtype=tf.int32, shape=[2]),
    event_shape_in= tf.placeholder(dtype=tf.int32, shape=[3]),
    validate_args=True)

  bijector.forward_event_shape(tf.TensorShape([5, 2, 3, 7]))
  # Either policy ==> (5, None, None)
  ```

  """

  def __init__(self, event_shape_out, event_shape_in=(-1,),
               validate_args=False, name=None):
    """Creates a `Reshape` bijector.

    Args:
      event_shape_out: An `int`-like vector-shaped `Tensor`
        representing the event shape of the transformed output.
      event_shape_in: An optional `int`-like vector-shape `Tensor`
        representing the event shape of the input. This is required in
        order to define inverse operations; the default of (-1,)
        assumes a vector-shaped input.
      validate_args: Python `bool` indicating whether arguments should
        be checked for correctness.
      name: Python `str`, name given to ops managed by this object.

    Raises:
      TypeError: if either `event_shape_in` or `event_shape_out` has
        non-integer `dtype`.
      ValueError: if either of `event_shape_in` or `event_shape_out`
       has non-vector shape (`rank > 1`), or if their sizes do not
       match.
    """
    with tf.name_scope(name or 'reshape'):
      event_shape_out = tf.convert_to_tensor(
          value=event_shape_out, name='event_shape_out', dtype_hint=tf.int32)
      event_shape_in = tf.convert_to_tensor(
          value=event_shape_in, name='event_shape_in', dtype_hint=tf.int32)

      forward_min_event_ndims_ = tensorshape_util.num_elements(
          event_shape_in.shape)
      if forward_min_event_ndims_ is None:
        raise NotImplementedError(
            '`event_shape_in` `size` must be statically known. For dynamic '
            'support, please contact `tfprobability@tensorflow.org`.')

      inverse_min_event_ndims_ = tensorshape_util.num_elements(
          event_shape_out.shape)
      if inverse_min_event_ndims_ is None:
        raise NotImplementedError(
            '`event_shape_out` `size` must be statically known. For dynamic '
            'support, please contact `tfprobability@tensorflow.org`.')

      assertions = []
      assertions.extend(_maybe_check_valid_shape(
          event_shape_out, validate_args))
      assertions.extend(_maybe_check_valid_shape(
          event_shape_in, validate_args))

      if assertions:
        with tf.control_dependencies(assertions):
          event_shape_in = tf.identity(
              event_shape_in, name='validated_event_shape_in')
          event_shape_out = tf.identity(
              event_shape_out, name='validated_event_shape_out')

      self._event_shape_in = event_shape_in
      self._event_shape_out = event_shape_out

      super(Reshape, self).__init__(
          forward_min_event_ndims=forward_min_event_ndims_,
          inverse_min_event_ndims=inverse_min_event_ndims_,
          is_constant_jacobian=True,
          validate_args=validate_args,
          name=name or 'reshape')

  def _forward(self, x):
    output_shape, output_tensorshape = _replace_event_shape_in_shape_tensor(
        tf.shape(input=x),
        self._event_shape_in,
        self._event_shape_out,
        self.validate_args)
    y = tf.reshape(x, output_shape)
    tensorshape_util.set_shape(y, output_tensorshape)
    return y

  def _inverse(self, y):
    output_shape, output_tensorshape = _replace_event_shape_in_shape_tensor(
        tf.shape(input=y),
        self._event_shape_out,
        self._event_shape_in,
        self.validate_args)
    x = tf.reshape(y, output_shape)
    tensorshape_util.set_shape(x, output_tensorshape)
    return x

  def _inverse_log_det_jacobian(self, y):
    return tf.constant(0., dtype=y.dtype)

  def _forward_log_det_jacobian(self, x):
    return tf.constant(0., dtype=x.dtype)

  def _forward_event_shape(self, input_shape):
    return _replace_event_shape_in_tensorshape(
        input_shape,
        self._event_shape_in,
        self._event_shape_out)[0]

  def _inverse_event_shape(self, output_shape):
    return _replace_event_shape_in_tensorshape(
        output_shape,
        self._event_shape_out,
        self._event_shape_in)[0]

  def _forward_event_shape_tensor(self, input_shape):
    return _replace_event_shape_in_shape_tensor(
        input_shape,
        self._event_shape_in,
        self._event_shape_out,
        self.validate_args)[0]

  def _inverse_event_shape_tensor(self, output_shape):
    return _replace_event_shape_in_shape_tensor(
        output_shape,
        self._event_shape_out,
        self._event_shape_in,
        self.validate_args)[0]

def _replace_event_shape_in_shape_tensor(
    input_shape, event_shape_in, event_shape_out, validate_args):
  """Replaces the rightmost dims in a `Tensor` representing a shape.

  Args:
    input_shape: a rank-1 `Tensor` of integers
    event_shape_in: the event shape expected to be present in rightmost dims
      of `shape_in`.
    event_shape_out: the event shape with which to replace `event_shape_in` in
      the rightmost dims of `input_shape`.
    validate_args: Python `bool` indicating whether arguments should
      be checked for correctness.

  Returns:
    output_shape: A rank-1 integer `Tensor` with the same contents as
      `input_shape` except for the event dims, which are replaced with
      `event_shape_out`.
  """
  output_tensorshape, is_validated = _replace_event_shape_in_tensorshape(
      tensorshape_util.constant_value_as_shape(input_shape),
      event_shape_in,
      event_shape_out)

  # TODO(b/124240153): Remove map(tf.identity, deps) once tf.function
  # correctly supports control_dependencies.
  validation_dependencies = (
      map(tf.identity, (event_shape_in, event_shape_out))
      if validate_args else ())

  if (tensorshape_util.is_fully_defined(output_tensorshape) and
      (is_validated or not validate_args)):
    with tf.control_dependencies(validation_dependencies):
      output_shape = tf.convert_to_tensor(
          value=output_tensorshape, name='output_shape', dtype_hint=tf.int32)
    return output_shape, output_tensorshape

  with tf.control_dependencies(validation_dependencies):
    event_shape_in_ndims = (
        tf.size(input=event_shape_in)
        if tensorshape_util.num_elements(event_shape_in.shape) is None else
        tensorshape_util.num_elements(event_shape_in.shape))
    input_non_event_shape, input_event_shape = tf.split(
        input_shape, num_or_size_splits=[-1, event_shape_in_ndims])

  additional_assertions = []
  if is_validated:
    pass
  elif validate_args:
    # Check that `input_event_shape` and `event_shape_in` are compatible in the
    # sense that they have equal entries in any position that isn't a `-1` in
    # `event_shape_in`. Note that our validations at construction time ensure
    # there is at most one such entry in `event_shape_in`.
    mask = event_shape_in >= 0
    explicit_input_event_shape = tf.boolean_mask(
        tensor=input_event_shape, mask=mask)
    explicit_event_shape_in = tf.boolean_mask(
        tensor=event_shape_in, mask=mask)
    additional_assertions.append(
        assert_util.assert_equal(
            explicit_input_event_shape,
            explicit_event_shape_in,
            message='Input `event_shape` does not match `event_shape_in`.'))
    # We don't explicitly additionally verify
    # `tf.size(input_shape) > tf.size(event_shape_in)` since `tf.split`
    # already makes this assertion.

  with tf.control_dependencies(additional_assertions):
    output_shape = tf.concat([input_non_event_shape, event_shape_out], axis=0,
                             name='output_shape')

  return output_shape, output_tensorshape

def _replace_event_shape_in_tensorshape(
    input_tensorshape, event_shape_in, event_shape_out):
  """Replaces the event shape dims of a `TensorShape`.

  Args:
    input_tensorshape: a `TensorShape` instance in which to attempt replacing
      event shape.
    event_shape_in: `Tensor` shape representing the event shape expected to
      be present in (rightmost dims of) `tensorshape_in`. Must be compatible
      with the rightmost dims of `tensorshape_in`.
    event_shape_out: `Tensor` shape representing the new event shape, i.e.,
      the replacement of `event_shape_in`,

  Returns:
    output_tensorshape: `TensorShape` with the rightmost `event_shape_in`
      replaced by `event_shape_out`. Might be partially defined, i.e.,
      `TensorShape(None)`.
    is_validated: Python `bool` indicating static validation happened.

  Raises:
    ValueError: if we can determine the event shape portion of
      `tensorshape_in` as well as `event_shape_in` both statically, and they
      are not compatible. "Compatible" here means that they are identical on
      any dims that are not -1 in `event_shape_in`.
  """
  event_shape_in_ndims = tensorshape_util.num_elements(event_shape_in.shape)
  if tensorshape_util.rank(
      input_tensorshape) is None or event_shape_in_ndims is None:
    return tf.TensorShape(None), False  # Not is_validated.

  input_non_event_ndims = tensorshape_util.rank(
      input_tensorshape) - event_shape_in_ndims
  if input_non_event_ndims < 0:
    raise ValueError(
        'Input has fewer ndims ({}) than event shape ndims ({}).'.format(
            tensorshape_util.rank(input_tensorshape), event_shape_in_ndims))

  input_non_event_tensorshape = input_tensorshape[:input_non_event_ndims]
  input_event_tensorshape = input_tensorshape[input_non_event_ndims:]

  # Check that `input_event_shape_` and `event_shape_in` are compatible in the
  # sense that they have equal entries in any position that isn't a `-1` in
  # `event_shape_in`. Note that our validations at construction time ensure
  # there is at most one such entry in `event_shape_in`.
  event_shape_in_ = tf.get_static_value(event_shape_in)
  is_validated = (
      tensorshape_util.is_fully_defined(input_event_tensorshape) and
      event_shape_in_ is not None)
  if is_validated:
    input_event_shape_ = np.int32(input_event_tensorshape)
    mask = event_shape_in_ >= 0
    explicit_input_event_shape_ = input_event_shape_[mask]
    explicit_event_shape_in_ = event_shape_in_[mask]
    if not all(explicit_input_event_shape_ == explicit_event_shape_in_):
      raise ValueError(
          'Input `event_shape` does not match `event_shape_in`. '
          '({} vs {}).'.format(input_event_shape_, event_shape_in_))

  event_tensorshape_out = tensorshape_util.constant_value_as_shape(
      event_shape_out)
  if tensorshape_util.rank(event_tensorshape_out) is None:
    output_tensorshape = tf.TensorShape(None)
  else:
    output_tensorshape = tensorshape_util.concatenate(
        input_non_event_tensorshape, event_tensorshape_out)

  return output_tensorshape, is_validated

def _maybe_check_valid_shape(shape, validate_args):
  """Check that a shape Tensor is int-type and otherwise sane."""
  if not dtype_util.is_integer(shape.dtype):
    raise TypeError('{} dtype ({}) should be `int`-like.'.format(
        shape, dtype_util.name(shape.dtype)))

  assertions = []

  message = '`{}` rank should be <= 1.'
  if tensorshape_util.rank(shape.shape) is not None:
    if tensorshape_util.rank(shape.shape) > 1:
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_less(
        tf.rank(shape), 2, message=message.format(shape)))

  shape_ = tf.get_static_value(shape)

  message = '`{}` elements must have at most one `-1`.'
  if shape_ is not None:
    if sum(shape_ == -1) > 1:
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_less(
        tf.reduce_sum(input_tensor=tf.cast(tf.equal(shape, -1), tf.int32)),
        2,
        message=message.format(shape)))

  message = '`{}` elements must be either positive integers or `-1`.'
  if shape_ is not None:
    if np.any(shape_ < -1):
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_greater(
        shape, -2, message=message.format(shape)))

  return assertions


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""The Beta distribution class."""

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

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

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

__all__ = [
    "Beta",
]

_beta_sample_note = """Note: `x` must have dtype `self.dtype` and be in
`[0, 1].` It must have a shape compatible with `self.batch_shape()`."""

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

  The Beta distribution is defined over the `(0, 1)` interval using parameters
  `concentration1` (aka "alpha") and `concentration0` (aka "beta").

  #### Mathematical Details

  The probability density function (pdf) is,

  ```none
  pdf(x; alpha, beta) = x**(alpha - 1) (1 - x)**(beta - 1) / Z
  Z = Gamma(alpha) Gamma(beta) / Gamma(alpha + beta)
  ```

  where:

  * `concentration1 = alpha`,
  * `concentration0 = beta`,
  * `Z` is the normalization constant, and,
  * `Gamma` is the [gamma function](
    https://en.wikipedia.org/wiki/Gamma_function).

  The concentration parameters represent mean total counts of a `1` or a `0`,
  i.e.,

  ```none
  concentration1 = alpha = mean * total_concentration
  concentration0 = beta  = (1. - mean) * total_concentration
  ```

  where `mean` in `(0, 1)` and `total_concentration` is a positive real number
  representing a mean `total_count = concentration1 + concentration0`.

  Distribution parameters are automatically broadcast in all functions; see
  examples for details.

  Warning: The samples can be zero due to finite precision.
  This happens more often when some of the concentrations are very small.
  Make sure to round the samples to `np.finfo(dtype).tiny` before computing the
  density.

  Samples of this distribution are reparameterized (pathwise differentiable).
  The derivatives are computed using the approach described in the paper

  [Michael Figurnov, Shakir Mohamed, Andriy Mnih.
  Implicit Reparameterization Gradients, 2018](https://arxiv.org/abs/1805.08498)

  #### Examples

  ```python
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  # Create a batch of three Beta distributions.
  alpha = [1, 2, 3]
  beta = [1, 2, 3]
  dist = tfd.Beta(alpha, beta)

  dist.sample([4, 5])  # Shape [4, 5, 3]

  # `x` has three batch entries, each with two samples.
  x = [[.1, .4, .5],
       [.2, .3, .5]]
  # Calculate the probability of each pair of samples under the corresponding
  # distribution in `dist`.
  dist.prob(x)         # Shape [2, 3]
  ```

  ```python
  # Create batch_shape=[2, 3] via parameter broadcast:
  alpha = [[1.], [2]]      # Shape [2, 1]
  beta = [3., 4, 5]        # Shape [3]
  dist = tfd.Beta(alpha, beta)

  # alpha broadcast as: [[1., 1, 1,],
  #                      [2, 2, 2]]
  # beta broadcast as:  [[3., 4, 5],
  #                      [3, 4, 5]]
  # batch_Shape [2, 3]
  dist.sample([4, 5])  # Shape [4, 5, 2, 3]

  x = [.2, .3, .5]
  # x will be broadcast as [[.2, .3, .5],
  #                         [.2, .3, .5]],
  # thus matching batch_shape [2, 3].
  dist.prob(x)         # Shape [2, 3]
  ```

  Compute the gradients of samples w.r.t. the parameters:

  ```python
  alpha = tf.constant(1.0)
  beta = tf.constant(2.0)
  dist = tfd.Beta(alpha, beta)
  samples = dist.sample(5)  # Shape [5]
  loss = tf.reduce_mean(tf.square(samples))  # Arbitrary loss function
  # Unbiased stochastic gradients of the loss function
  grads = tf.gradients(loss, [alpha, beta])
  ```

  """

  def __init__(self,
               concentration1,
               concentration0,
               validate_args=False,
               allow_nan_stats=True,
               name="Beta"):
    """Initialize a batch of Beta distributions.

    Args:
      concentration1: Positive floating-point `Tensor` indicating mean
        number of successes; aka "alpha". Implies `self.dtype` and
        `self.batch_shape`, i.e.,
        `concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
      concentration0: Positive floating-point `Tensor` indicating mean
        number of failures; aka "beta". Otherwise has same semantics as
        `concentration1`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([concentration1, concentration0],
                                      tf.float32)
      self._concentration1 = self._maybe_assert_valid_concentration(
          tf.convert_to_tensor(
              value=concentration1, name="concentration1", dtype=dtype),
          validate_args)
      self._concentration0 = self._maybe_assert_valid_concentration(
          tf.convert_to_tensor(
              value=concentration0, name="concentration0", dtype=dtype),
          validate_args)
      self._total_concentration = self._concentration1 + self._concentration0
    super(Beta, self).__init__(
        dtype=dtype,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        parameters=parameters,
        graph_parents=[
            self._concentration1, self._concentration0,
            self._total_concentration
        ],
        name=name)

  @staticmethod
  def _param_shapes(sample_shape):
    s = tf.convert_to_tensor(value=sample_shape, dtype=tf.int32)
    return dict(concentration1=s, concentration0=s)

  @classmethod
  def _params_event_ndims(cls):
    return dict(concentration1=0, concentration0=0)

  @property
  def concentration1(self):
    """Concentration parameter associated with a `1` outcome."""
    return self._concentration1

  @property
  def concentration0(self):
    """Concentration parameter associated with a `0` outcome."""
    return self._concentration0

  @property
  def total_concentration(self):
    """Sum of concentration parameters."""
    return self._total_concentration

  def _batch_shape_tensor(self):
    return tf.shape(input=self.total_concentration)

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

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

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

  def _sample_n(self, n, seed=None):
    seed = seed_stream.SeedStream(seed, "beta")
    expanded_concentration1 = tf.ones_like(
        self.total_concentration, dtype=self.dtype) * self.concentration1
    expanded_concentration0 = tf.ones_like(
        self.total_concentration, dtype=self.dtype) * self.concentration0
    gamma1_sample = tf.random.gamma(
        shape=[n], alpha=expanded_concentration1, dtype=self.dtype, seed=seed())
    gamma2_sample = tf.random.gamma(
        shape=[n], alpha=expanded_concentration0, dtype=self.dtype, seed=seed())
    beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
    return beta_sample

  @distribution_util.AppendDocstring(_beta_sample_note)
  def _log_prob(self, x):
    return self._log_unnormalized_prob(x) - self._log_normalization()

  @distribution_util.AppendDocstring(_beta_sample_note)
  def _prob(self, x):
    return tf.exp(self._log_prob(x))

  @distribution_util.AppendDocstring(_beta_sample_note)
  def _log_cdf(self, x):
    return tf.math.log(self._cdf(x))

  @distribution_util.AppendDocstring(_beta_sample_note)
  def _cdf(self, x):
    return tf.math.betainc(self.concentration1, self.concentration0, x)

  def _log_unnormalized_prob(self, x):
    x = self._maybe_assert_valid_sample(x)
    return (tf.math.xlogy(self.concentration1 - 1., x) +
            (self.concentration0 - 1.) * tf.math.log1p(-x))

  def _log_normalization(self):
    return (tf.math.lgamma(self.concentration1) +
            tf.math.lgamma(self.concentration0) -
            tf.math.lgamma(self.total_concentration))

  def _entropy(self):
    return (self._log_normalization() -
            (self.concentration1 - 1.) * tf.math.digamma(self.concentration1) -
            (self.concentration0 - 1.) * tf.math.digamma(self.concentration0) +
            ((self.total_concentration - 2.) *
             tf.math.digamma(self.total_concentration)))

  def _mean(self):
    return self._concentration1 / self._total_concentration

  def _variance(self):
    return self._mean() * (1. - self._mean()) / (1. + self.total_concentration)

  @distribution_util.AppendDocstring(
      """Note: The mode is undefined when `concentration1 <= 1` or
      `concentration0 <= 1`. If `self.allow_nan_stats` is `True`, `NaN`
      is used for undefined modes. If `self.allow_nan_stats` is `False` an
      exception is raised when one or more modes are undefined.""")
  def _mode(self):
    mode = (self.concentration1 - 1.) / (self.total_concentration - 2.)
    if self.allow_nan_stats:
      nan = tf.fill(
          self.batch_shape_tensor(),
          dtype_util.as_numpy_dtype(self.dtype)(np.nan),
          name="nan")
      is_defined = tf.logical_and(self.concentration1 > 1.,
                                  self.concentration0 > 1.)
      return tf.where(is_defined, mode, nan)
    return distribution_util.with_dependencies([
        assert_util.assert_less(
            tf.ones([], dtype=self.dtype),
            self.concentration1,
            message="Mode undefined for concentration1 <= 1."),
        assert_util.assert_less(
            tf.ones([], dtype=self.dtype),
            self.concentration0,
            message="Mode undefined for concentration0 <= 1.")
    ], mode)

  def _maybe_assert_valid_concentration(self, concentration, validate_args):
    """Checks the validity of a concentration parameter."""
    if not validate_args:
      return concentration
    return distribution_util.with_dependencies([
        assert_util.assert_positive(
            concentration, message="Concentration parameter must be positive."),
    ], concentration)

  def _maybe_assert_valid_sample(self, x):
    """Checks the validity of a sample."""
    if not self.validate_args:
      return x
    return distribution_util.with_dependencies([
        assert_util.assert_positive(x, message="sample must be positive"),
        assert_util.assert_less(x, 1., message="sample must be less than `1`."),
    ], x)

@kullback_leibler.RegisterKL(Beta, Beta)
def _kl_beta_beta(d1, d2, name=None):
  """Calculate the batchwise KL divergence KL(d1 || d2) with d1 and d2 Beta.

  Args:
    d1: instance of a Beta distribution object.
    d2: instance of a Beta distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_beta_beta".

  Returns:
    Batchwise KL(d1 || d2)
  """
  def delta(fn, is_property=True):
    fn1 = getattr(d1, fn)
    fn2 = getattr(d2, fn)
    return (fn2 - fn1) if is_property else (fn2() - fn1())

  with tf.name_scope(name or "kl_beta_beta"):
    return (delta("_log_normalization", is_property=False) -
            tf.math.digamma(d1.concentration1) * delta("concentration1") -
            tf.math.digamma(d1.concentration0) * delta("concentration0") +
            (tf.math.digamma(d1.total_concentration) *
             delta("total_concentration")))


# 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.
# ============================================================================
"""Common functions for BFGS and L-BFGS algorithm."""

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

import collections
import numpy as np
import tensorflow as tf

from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import prefer_static
from tensorflow_probability.python.optimizer import linesearch

# A namedtuple to hold the point at which a line function is evaluated, the
# value of the function, directional derivative, and full gradient evaluated
# evaluated at that point. To be used with the linesearch method.
ValueAndGradient = collections.namedtuple('ValueAndGradient',
                                          ['x', 'f', 'df', 'full_gradient'])

def converged_any(converged, failed):
  """Condition to stop when any batch member converges, or all have failed."""
  return (tf.reduce_any(input_tensor=converged) |
          tf.reduce_all(input_tensor=failed))

def converged_all(converged, failed):
  """Condition to stop when all batch members have converged or failed."""
  return tf.reduce_all(input_tensor=converged | failed)

def get_initial_state_args(value_and_gradients_function,
                           initial_position,
                           grad_tolerance,
                           control_inputs=None):
  """Returns a dictionary to populate the initial state of the search procedure.

  Performs an initial convergence check and the first evaluation of the
  objective function.

  Args:
    value_and_gradients_function: A Python callable that accepts a tensor and
      returns a tuple of two tensors: the objective function value and its
      derivative.
    initial_position: The starting point of the search procedure.
    grad_tolerance: The gradient tolerance for the procedure.
    control_inputs: Optional ops used to assert the validity of inputs, these
      are added as control dependencies to execute before the objective
      function is evaluated for the first time.

  Returns:
    An dictionary with values for the following keys:
      converged: True if the convergence check finds that the initial position
        is already an argmin of the objective function.
      failed: Initialized to False.
      num_objective_evaluations: Initialized to 1.
      position: Initialized to the initial position.
      objective_value: Initialized to the value of the objective function at
        the initial position.
      objective_gradient: Initialized to the gradient of the objective
        function at the initial position.
  """
  if control_inputs:
    with tf.control_dependencies(control_inputs):
      f0, df0 = value_and_gradients_function(initial_position)
  else:
    f0, df0 = value_and_gradients_function(initial_position)
  converged = norm(df0, dims=1) < grad_tolerance
  return dict(
      converged=converged,
      failed=tf.zeros_like(converged),  # i.e. False.
      num_iterations=tf.convert_to_tensor(value=0),
      num_objective_evaluations=tf.convert_to_tensor(value=1),
      position=initial_position,
      objective_value=f0,
      objective_gradient=df0)

def line_search_step(state, value_and_gradients_function, search_direction,
                     grad_tolerance, f_relative_tolerance, x_tolerance,
                     stopping_condition):
  """Performs the line search step of the BFGS search procedure.

  Uses hager_zhang line search procedure to compute a suitable step size
  to advance the current `state.position` along the given `search_direction`.
  Also, if the line search is successful, updates the `state.position` by
  taking the corresponding step.

  Args:
    state: A namedtuple instance holding values for the current state of the
      search procedure. The state must include the fields: `position`,
      `objective_value`, `objective_gradient`, `num_iterations`,
      `num_objective_evaluations`, `converged` and `failed`.
    value_and_gradients_function: A Python callable that accepts a point as a
      real `Tensor` of shape `[..., n]` and returns a tuple of two tensors of
      the same dtype: the objective function value, a real `Tensor` of shape
      `[...]`, and its derivative, another real `Tensor` of shape `[..., n]`.
    search_direction: A real `Tensor` of shape `[..., n]`. The direction along
      which to perform line search.
    grad_tolerance: Scalar `Tensor` of real dtype. Specifies the gradient
      tolerance for the procedure.
    f_relative_tolerance: Scalar `Tensor` of real dtype. Specifies the
      tolerance for the relative change in the objective value.
    x_tolerance: Scalar `Tensor` of real dtype. Specifies the tolerance for the
      change in the position.
    stopping_condition: A Python function that takes as input two Boolean
      tensors of shape `[...]`, and returns a Boolean scalar tensor. The input
      tensors are `converged` and `failed`, indicating the current status of
      each respective batch member; the return value states whether the
      algorithm should stop.
  Returns:
    A copy of the input state with the following fields updated:
      converged: a Boolean `Tensor` of shape `[...]` indicating whether the
        convergence criteria has been met.
      failed: a Boolean `Tensor` of shape `[...]` indicating whether the line
        search procedure failed to converge, or if either the updated gradient
        or objective function are no longer finite.
      num_iterations: Increased by 1.
      num_objective_evaluations: Increased by the number of times that the
        objective function got evaluated.
      position, objective_value, objective_gradient: If line search succeeded,
        updated by computing the new position and evaluating the objective
        function at that position.
  """
  line_search_value_grad_func = _restrict_along_direction(
      value_and_gradients_function, state.position, search_direction)
  derivative_at_start_pt = tf.reduce_sum(
      input_tensor=state.objective_gradient * search_direction, axis=-1)
  val_0 = ValueAndGradient(x=_broadcast(0, state.position),
                           f=state.objective_value,
                           df=derivative_at_start_pt,
                           full_gradient=state.objective_gradient)
  inactive = state.failed | state.converged
  ls_result = linesearch.hager_zhang(
      line_search_value_grad_func,
      initial_step_size=_broadcast(1, state.position),
      value_at_zero=val_0,
      converged=inactive)  # No search needed for these.

  state_after_ls = update_fields(
      state,
      failed=state.failed | ~ls_result.converged,
      num_iterations=state.num_iterations + 1,
      num_objective_evaluations=(
          state.num_objective_evaluations + ls_result.func_evals))

  def _do_update_position():
    # For inactive batch members `left.x` is zero. However, their
    # `search_direction` might also be undefined, so we can't rely on
    # multiplication by zero to produce a `position_delta` of zero.
    position_delta = tf.where(
        inactive,
        tf.zeros_like(search_direction),
        search_direction * tf.expand_dims(ls_result.left.x, axis=-1))
    return _update_position(
        state_after_ls,
        position_delta,
        ls_result.left.f,
        ls_result.left.full_gradient,
        grad_tolerance, f_relative_tolerance, x_tolerance)

  return prefer_static.cond(
      stopping_condition(state.converged, state.failed),
      true_fn=lambda: state_after_ls,
      false_fn=_do_update_position)

def update_fields(state, **kwargs):
  """Copies the argument and overrides some of its fields.

  Args:
    state: A `collections.namedtuple` instance.
    **kwargs: Other named arguments represent fields in the tuple to override
      with new values.

  Returns:
    A namedtuple, of the same class as the input argument, with the updated
    fields.

  Raises:
    ValueError if the supplied kwargs contain fields not present in the
    input argument.
  """
  return state._replace(**kwargs)

def _restrict_along_direction(value_and_gradients_function,
                              position,
                              direction):
  """Restricts a function in n-dimensions to a given direction.

  Suppose f: R^n -> R. Then given a point x0 and a vector p0 in R^n, the
  restriction of the function along that direction is defined by:

  ```None
  g(t) = f(x0 + t * p0)
  ```

  This function performs this restriction on the given function. In addition, it
  also computes the gradient of the restricted function along the restriction
  direction. This is equivalent to computing `dg/dt` in the definition above.

  Args:
    value_and_gradients_function: Callable accepting a single real `Tensor`
      argument of shape `[..., n]` and returning a tuple of a real `Tensor` of
      shape `[...]` and a real `Tensor` of shape `[..., n]`. The multivariate
      function whose restriction is to be computed. The output values of the
      callable are the function value and the gradients at the input argument.
    position: `Tensor` of real dtype and shape consumable by
      `value_and_gradients_function`. Corresponds to `x0` in the definition
      above.
    direction: `Tensor` of the same dtype and shape as `position`. The direction
      along which to restrict the function. Note that the direction need not
      be a unit vector.

  Returns:
    restricted_value_and_gradients_func: A callable accepting a tensor of shape
      broadcastable to `[...]` and same dtype as `position` and returning a
      namedtuple of `Tensors`. The input tensor is the parameter along the
      direction labelled `t` above. The return value contains fields:
        x: A real `Tensor` of shape `[...]`. The input value `t` where the line
          function was evaluated, after any necessary broadcasting.
        f: A real `Tensor` of shape `[...]` containing the value of the
          function at the point `position + t * direction`.
        df: A real `Tensor` of shape `[...]` containing the derivative at
          `position + t * direction`.
        full_gradient: A real `Tensor` of shape `[..., n]`, the full gradient
          of the original `value_and_gradients_function`.
  """
  def _restricted_func(t):
    t = _broadcast(t, position)
    pt = position + tf.expand_dims(t, axis=-1) * direction
    objective_value, gradient = value_and_gradients_function(pt)
    return ValueAndGradient(
        x=t,
        f=objective_value,
        df=tf.reduce_sum(input_tensor=gradient * direction, axis=-1),
        full_gradient=gradient)

  return _restricted_func

def _update_position(state,
                     position_delta,
                     next_objective,
                     next_gradient,
                     grad_tolerance,
                     f_relative_tolerance,
                     x_tolerance):
  """Updates the state advancing its position by a given position_delta."""
  failed = state.failed | ~tf.math.is_finite(next_objective) | ~tf.reduce_all(
      input_tensor=tf.math.is_finite(next_gradient), axis=-1)

  next_position = state.position + position_delta
  converged = ~failed & _check_convergence(state.position,
                                           next_position,
                                           state.objective_value,
                                           next_objective,
                                           next_gradient,
                                           grad_tolerance,
                                           f_relative_tolerance,
                                           x_tolerance)
  return update_fields(
      state,
      converged=state.converged | converged,
      failed=failed,
      position=next_position,
      objective_value=next_objective,
      objective_gradient=next_gradient)

def norm(value, dims, order=None):
  """Compute the norm of the given (possibly batched) value.

  Args:
    value: A `Tensor` of real dtype.
    dims: An Python integer with the number of non-batching dimensions in the
      value, i.e. `dims=0` (scalars), `dims=1` (vectors), `dims=2` (matrices).
    order: Order of the norm, defaults to `np.inf`.
  """
  if dims == 0:
    return tf.math.abs(value)
  elif dims == 1:
    axis = -1
  elif dims == 2:
    axis = [-1, -2]
  else:
    ValueError(dims)
  if order is None:
    order = np.inf
  return tf.norm(tensor=value, axis=axis, ord=order)

def _check_convergence(current_position,
                       next_position,
                       current_objective,
                       next_objective,
                       next_gradient,
                       grad_tolerance,
                       f_relative_tolerance,
                       x_tolerance):
  """Checks if the algorithm satisfies the convergence criteria."""
  grad_converged = norm(next_gradient, dims=1) <= grad_tolerance
  x_converged = norm(next_position - current_position, dims=1) <= x_tolerance
  f_converged = (norm(next_objective - current_objective, dims=0) <=
                 f_relative_tolerance * current_objective)
  return grad_converged | x_converged | f_converged

def _broadcast(value, target):
  """Broadcast a value to match the batching dimensions of a target.

  If necessary the value is converted into a tensor. Both value and target
  should be of the same dtype.

  Args:
    value: A value to broadcast.
    target: A `Tensor` of shape [b1, ..., bn, d].

  Returns:
    A `Tensor` of shape [b1, ..., bn] and same dtype as the target.
  """
  return tf.broadcast_to(
      tf.convert_to_tensor(value=value, dtype=target.dtype),
      distribution_util.prefer_static_shape(target)[:-1])


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""The Kumaraswamy distribution class."""

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

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

from tensorflow_probability.python.bijectors import kumaraswamy as kumaraswamy_bijector
from tensorflow_probability.python.distributions import transformed_distribution
from tensorflow_probability.python.distributions import uniform
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util

__all__ = [
    "Kumaraswamy",
]

_kumaraswamy_sample_note = """Note: `x` must have dtype `self.dtype` and be in
`[0, 1].` It must have a shape compatible with `self.batch_shape()`."""

def _harmonic_number(x):
  """Compute the harmonic number from its analytic continuation.

  Derivation from [here](
  https://en.wikipedia.org/wiki/Digamma_function#Relation_to_harmonic_numbers)
  and [Euler's constant](
  https://en.wikipedia.org/wiki/Euler%E2%80%93Mascheroni_constant).

  Args:
    x: input float.

  Returns:
    z: The analytic continuation of the harmonic number for the input.
  """
  one = tf.ones([], dtype=x.dtype)
  return tf.math.digamma(x + one) - tf.math.digamma(one)

class Kumaraswamy(transformed_distribution.TransformedDistribution):
  """Kumaraswamy distribution.

  The Kumaraswamy distribution is defined over the `(0, 1)` interval using
  parameters
  `concentration1` (aka "alpha") and `concentration0` (aka "beta").  It has a
  shape similar to the Beta distribution, but is easier to reparameterize.

  #### Mathematical Details

  The probability density function (pdf) is,

  ```none
  pdf(x; alpha, beta) = alpha * beta * x**(alpha - 1) * (1 - x**alpha)**(beta -
  1)
  ```

  where:

  * `concentration1 = alpha`,
  * `concentration0 = beta`,

  Distribution parameters are automatically broadcast in all functions; see
  examples for details.

  #### Examples

  ```python
  # Create a batch of three Kumaraswamy distributions.
  alpha = [1, 2, 3]
  beta = [1, 2, 3]
  dist = Kumaraswamy(alpha, beta)

  dist.sample([4, 5])  # Shape [4, 5, 3]

  # `x` has three batch entries, each with two samples.
  x = [[.1, .4, .5],
       [.2, .3, .5]]
  # Calculate the probability of each pair of samples under the corresponding
  # distribution in `dist`.
  dist.prob(x)         # Shape [2, 3]
  ```

  ```python
  # Create batch_shape=[2, 3] via parameter broadcast:
  alpha = [[1.], [2]]      # Shape [2, 1]
  beta = [3., 4, 5]        # Shape [3]
  dist = Kumaraswamy(alpha, beta)

  # alpha broadcast as: [[1., 1, 1,],
  #                      [2, 2, 2]]
  # beta broadcast as:  [[3., 4, 5],
  #                      [3, 4, 5]]
  # batch_Shape [2, 3]
  dist.sample([4, 5])  # Shape [4, 5, 2, 3]

  x = [.2, .3, .5]
  # x will be broadcast as [[.2, .3, .5],
  #                         [.2, .3, .5]],
  # thus matching batch_shape [2, 3].
  dist.prob(x)         # Shape [2, 3]
  ```

  """

  def __init__(self,
               concentration1=1.,
               concentration0=1.,
               validate_args=False,
               allow_nan_stats=True,
               name="Kumaraswamy"):
    """Initialize a batch of Kumaraswamy distributions.

    Args:
      concentration1: Positive floating-point `Tensor` indicating mean
        number of successes; aka "alpha". Implies `self.dtype` and
        `self.batch_shape`, i.e.,
        `concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
      concentration0: Positive floating-point `Tensor` indicating mean
        number of failures; aka "beta". Otherwise has same semantics as
        `concentration1`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([concentration1, concentration0],
                                      tf.float32)
      concentration1 = tf.convert_to_tensor(
          value=concentration1, name="concentration1", dtype=dtype)
      concentration0 = tf.convert_to_tensor(
          value=concentration0, name="concentration0", dtype=dtype)
    super(Kumaraswamy, self).__init__(
        distribution=uniform.Uniform(
            low=tf.zeros([], dtype=concentration1.dtype),
            high=tf.ones([], dtype=concentration1.dtype),
            allow_nan_stats=allow_nan_stats),
        bijector=kumaraswamy_bijector.Kumaraswamy(
            concentration1=concentration1,
            concentration0=concentration0,
            validate_args=validate_args),
        batch_shape=distribution_util.get_broadcast_shape(
            concentration1, concentration0),
        parameters=parameters,
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(concentration1=0, concentration0=0)

  @property
  def concentration1(self):
    """Concentration parameter associated with a `1` outcome."""
    return self.bijector.concentration1

  @property
  def concentration0(self):
    """Concentration parameter associated with a `0` outcome."""
    return self.bijector.concentration0

  def _entropy(self):
    a = self.concentration1
    b = self.concentration0
    return ((1 - 1. / b) + (1 - 1. / a) * _harmonic_number(b) -
            tf.math.log(a) - tf.math.log(b))

  def _moment(self, n):
    """Compute the n'th (uncentered) moment."""
    total_concentration = self.concentration1 + self.concentration0
    expanded_concentration1 = tf.ones_like(
        total_concentration, dtype=self.dtype) * self.concentration1
    expanded_concentration0 = tf.ones_like(
        total_concentration, dtype=self.dtype) * self.concentration0
    beta_arg0 = 1 + n / expanded_concentration1
    beta_arg = tf.stack([beta_arg0, expanded_concentration0], -1)
    log_moment = tf.math.log(expanded_concentration0) + tf.math.lbeta(beta_arg)
    return tf.exp(log_moment)

  def _mean(self):
    return self._moment(1)

  def _variance(self):
    # TODO(b/72696533): Investigate a more numerically stable version.
    return self._moment(2) - tf.square(self._moment(1))

  @distribution_util.AppendDocstring(
      """Note: The mode is undefined when `concentration1 <= 1` or
      `concentration0 <= 1`. If `self.allow_nan_stats` is `True`, `NaN`
      is used for undefined modes. If `self.allow_nan_stats` is `False` an
      exception is raised when one or more modes are undefined.""")
  def _mode(self):
    a = self.concentration1
    b = self.concentration0
    mode = ((a - 1) / (a * b - 1))**(1. / a)
    if self.allow_nan_stats:
      nan = tf.fill(
          self.batch_shape_tensor(),
          dtype_util.as_numpy_dtype(self.dtype)(np.nan),
          name="nan")
      is_defined = (self.concentration1 > 1.) & (self.concentration0 > 1.)
      return tf.where(is_defined, mode, nan)

    return distribution_util.with_dependencies([
        assert_util.assert_less(
            tf.ones([], dtype=self.concentration1.dtype),
            self.concentration1,
            message="Mode undefined for concentration1 <= 1."),
        assert_util.assert_less(
            tf.ones([], dtype=self.concentration0.dtype),
            self.concentration0,
            message="Mode undefined for concentration0 <= 1.")
    ], mode)


# 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.
# ============================================================================
"""SimpleStepSizeAdaptation TransitionKernel."""

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

import collections

import numpy as np
import tensorflow as tf

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

@mcmc_util.make_innermost_setter
def _hmc_like_step_size_setter_fn(kernel_results, new_step_size):
  return kernel_results._replace(
      accepted_results=kernel_results.accepted_results._replace(
          step_size=new_step_size))

@mcmc_util.make_innermost_getter
def _hmc_like_step_size_getter_fn(kernel_results):
  return kernel_results.accepted_results.step_size

@mcmc_util.make_innermost_getter
def _hmc_like_log_accept_prob_getter_fn(kernel_results):
  log_accept_ratio = kernel_results.log_accept_ratio
  return tf.minimum(tf.constant(0., log_accept_ratio.dtype), log_accept_ratio)

def _reduce_logmeanexp(value, dims, keepdims=False):
  # This is intentionally numerically imprecise for simplicity. For the purposes
  # of computing the mean acceptance probability this is more than sufficient.
  return tf.math.log(
      tf.reduce_mean(input_tensor=tf.exp(value), axis=dims, keepdims=keepdims))

def _get_differing_dims(a, b):
  # Get the indices of dimensions where shapes of `a` and `b` differ.
  # `a` is allowed to have fewer dimensions than `b`.
  if a.shape.is_fully_defined() and b.shape.is_fully_defined():
    a_shape = np.array(a.shape.as_list())
    b_shape = np.array(b.shape.as_list())
    return np.where(a_shape != b_shape[:len(a_shape)])[0]
  else:
    return tf.where(
        tf.not_equal(tf.shape(input=a),
                     tf.shape(input=b)[:tf.rank(a)]))[:, 0]

class SimpleStepSizeAdaptationResults(
    collections.namedtuple(
        'SimpleStepSizeAdaptationResults',
        'inner_results, target_accept_prob, adaptation_rate, step, '
        'new_step_size')):
  """Results of the SimpleStepSizeAdaptation TransitionKernel.

  Attributes:
    inner_results: Results of the inner kernel.
    target_accept_prob: Floating point scalar `Tensor`. Target accept
      probability.
    adaptation_rate: Floating point scalar `Tensor`. Fraction by which to adjust
      the step size during each step.
    step: Int32 scalar `Tensor`. The current step number as perceived by this
      kernel. Increases by 1 for every call to `one_step`.
    new_step_size:  Floating point scalar `Tensor` or a list thereof (one for
      each `state_part`). Step size that will be passed to the inner kernel
      during the next step.
  """
  __slots__ = ()

class SimpleStepSizeAdaptation(kernel_base.TransitionKernel):
  """Adapts the inner kernel's `step_size` based on `log_accept_prob`.

  The simple policy multiplicatively increases or decreases the `step_size` of
  the inner kernel based on the value of `log_accept_prob`. It is based on
  [equation 19 of Andrieu and Thoms (2008)][1]. Given enough steps and small
  enough `adaptation_rate` the median of the distribution of the acceptance
  probability will converge to the `target_accept_prob`. A good target
  acceptance probability depends on the inner kernel. If this kernel is
  `HamiltonianMonteCarlo`, then 0.6-0.9 is a good range to aim for. For
  `RandomWalkMetropolis` this should be closer to 0.25. See the individual
  kernels' docstrings for guidance.

  In general, adaptation prevents the chain from reaching a stationary
  distribution, so obtaining consistent samples requires `num_adaptation_steps`
  be set to a value [somewhat smaller][2] than the number of burnin steps.
  However, it may sometimes be helpful to set `num_adaptation_steps` to a larger
  value during development in order to inspect the behavior of the chain during
  adaptation.

  The step size is assumed to broadcast with the chain state, potentially having
  leading dimensions corresponding to multiple chains. When there are fewer of
  those leading dimensions than there are chain dimensions, the corresponding
  dimensions in the `log_accept_prob` are averaged (in the direct space, rather
  than the log space) before being used to adjust the step size. This means that
  this kernel can do both cross-chain adaptation, or per-chain step size
  adaptation, depending on the shape of the step size.

  For example, if your problem has a state with shape `[S]`, your chain state
  has shape `[C0, C1, Y]` (meaning that there are `C0 * C1` total chains) and
  `log_accept_prob` has shape `[C0, C1]` (one acceptance probability per chain),
  then depending on the shape of the step size, the following will happen:

  - Step size has shape [], [S] or [1], the `log_accept_prob` will be averaged
    across its `C0` and `C1` dimensions. This means that you will learn a shared
    step size based on the mean acceptance probability across all chains. This
    can be useful if you don't have a lot of steps to adapt and want to average
    away the noise.

  - Step size has shape [C1, 1] or [C1, S], the `log_accept_prob` will be
    averaged across its `C0` dimension. This means that you will learn a shared
    step size based on the mean acceptance probability across chains that share
    the coordinate across the `C1` dimension. This can be useful when the `C1`
    dimension indexes different distributions, while `C0` indexes replicas of a
    single distribution, all sampled in parallel.

  - Step size has shape [C0, C1, 1] or [C0, C1, S], then no averaging will
    happen. This means that each chain will learn its own step size. This can be
    useful when all chains are sampling from different distributions. Even when
    all chains are for the same distribution, this can help during the initial
    warmup period.

  - Step size has shape [C0, 1, 1] or [C0, 1, S], the `log_accept_prob` will be
    averaged across its `C1` dimension. This means that you will learn a shared
    step size based on the mean acceptance probability across chains that share
    the coordinate across the `C0` dimension. This can be useful when the `C0`
    dimension indexes different distributions, while `C1` indexes replicas of a
    single distribution, all sampled in parallel.

  #### Examples

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

  target_log_prob_fn = tfd.Normal(loc=0., scale=1.).log_prob
  num_burnin_steps = 500
  num_results = 500
  num_chains = 64
  step_size = 0.1
  # Or, if you want per-chain step size:
  # step_size = tf.fill([num_chains], step_size)

  kernel = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=target_log_prob_fn,
      num_leapfrog_steps=2,
      step_size=step_size)
  kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=kernel, num_adaptation_steps=int(num_burnin_steps * 0.8))

  # The chain will be stepped for num_results + num_burnin_steps, adapting for
  # the first num_adaptation_steps.
  samples, [step_size, log_accept_ratio] = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=tf.zeros(num_chains),
      kernel=kernel,
      trace_fn=lambda _, pkr: [pkr.inner_results.accepted_results.step_size,
                               pkr.inner_results.log_accept_ratio])

  # ~0.75
  p_accept = tf.reduce_mean(tf.exp(tf.minimum(log_accept_ratio, 0.)))
  ```

  #### References

  [1]: Andrieu, Christophe, Thoms, Johannes. A tutorial on adaptive MCMC.
       _Statistics and Computing_, 2008.
       https://people.eecs.berkeley.edu/~jordan/sail/readings/andrieu-thoms.pdf

  [2]:
  http://andrewgelman.com/2017/12/15/burn-vs-warm-iterative-simulation-algorithms/#comment-627745

  """

  def __init__(self,
               inner_kernel,
               num_adaptation_steps,
               target_accept_prob=0.75,
               adaptation_rate=0.01,
               step_size_setter_fn=_hmc_like_step_size_setter_fn,
               step_size_getter_fn=_hmc_like_step_size_getter_fn,
               log_accept_prob_getter_fn=_hmc_like_log_accept_prob_getter_fn,
               validate_args=False,
               name=None):
    """Creates the step size adaptation kernel.

    The default setter_fn and the getter_fn callbacks assume that the inner
    kernel produces kernel results structurally the same as the
    `HamiltonianMonteCarlo` kernel.

    Args:
      inner_kernel: `TransitionKernel`-like object.
      num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
        during which to adjust the step size. This may be greater, less than, or
        equal to the number of burnin steps.
      target_accept_prob: A floating point `Tensor` representing desired
        acceptance probability. Must be a positive number less than 1. This can
        either be a scalar, or have shape [num_chains]. Default value: `0.75`
        (the [center of asymptotically optimal rate for HMC][1]).
      adaptation_rate: `Tensor` representing amount to scale the current
        `step_size`.
      step_size_setter_fn: A callable with the signature
        `(kernel_results, new_step_size) -> new_kernel_results` where
        `kernel_results` are the results of the `inner_kernel`, `new_step_size`
        is a `Tensor` or a nested collection of `Tensor`s with the same
        structure as returned by the `step_size_getter_fn`, and
        `new_kernel_results` are a copy of `kernel_results` with the step
        size(s) set.
      step_size_getter_fn: A callable with the signature
        `(kernel_results) -> step_size` where `kernel_results` are the results
        of the `inner_kernel`, and `step_size` is a floating point `Tensor` or a
        nested collection of such `Tensor`s.
      log_accept_prob_getter_fn: A callable with the signature
        `(kernel_results) -> log_accept_prob` where `kernel_results` are the
        results of the `inner_kernel`, and `log_accept_prob` is a floating point
        `Tensor`. `log_accept_prob` can either be a scalar, or have shape
        [num_chains]. If it's the latter, `step_size` should also have the same
        leading dimension.
      validate_args: Python `bool`. When `True` kernel parameters are checked
        for validity. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class. Default:
        'simple_step_size_adaptation'.

    #### References

    [1]: Betancourt, M. J., Byrne, S., & Girolami, M. (2014). _Optimizing The
         Integrator Step Size for Hamiltonian Monte Carlo_.
         http://arxiv.org/abs/1411.6669
    """

    inner_kernel = mcmc_util.enable_store_parameters_in_results(inner_kernel)

    with tf.compat.v1.name_scope(
        mcmc_util.make_name(name, 'simple_step_size_adaptation', '__init__'),
        values=[target_accept_prob, adaptation_rate,
                num_adaptation_steps]) as name:
      dtype = dtype_util.common_dtype([target_accept_prob, adaptation_rate],
                                      tf.float32)
      target_accept_prob = tf.convert_to_tensor(
          value=target_accept_prob, dtype=dtype, name='target_accept_prob')
      adaptation_rate = tf.convert_to_tensor(
          value=adaptation_rate, dtype=dtype, name='adaptation_rate')
      num_adaptation_steps = tf.convert_to_tensor(
          value=num_adaptation_steps,
          dtype=tf.int32,
          name='num_adaptation_steps')

      target_accept_prob = _maybe_validate_target_accept_prob(
          target_accept_prob, validate_args)

    self._parameters = dict(
        inner_kernel=inner_kernel,
        num_adaptation_steps=num_adaptation_steps,
        target_accept_prob=target_accept_prob,
        adaptation_rate=adaptation_rate,
        step_size_setter_fn=step_size_setter_fn,
        step_size_getter_fn=step_size_getter_fn,
        log_accept_prob_getter_fn=log_accept_prob_getter_fn,
        name=name,
    )

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

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

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

  def step_size_setter_fn(self, kernel_results, new_step_size):
    return self._parameters['step_size_setter_fn'](kernel_results,
                                                   new_step_size)

  def step_size_getter_fn(self, kernel_results):
    return self._parameters['step_size_getter_fn'](kernel_results)

  def log_accept_prob_getter_fn(self, kernel_results):
    return self._parameters['log_accept_prob_getter_fn'](kernel_results)

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

  def one_step(self, current_state, previous_kernel_results):
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'simple_step_size_adaptation',
                                 'one_step'),
        values=[current_state, previous_kernel_results]):

      # Set the step_size.
      inner_results = self.step_size_setter_fn(
          previous_kernel_results.inner_results,
          previous_kernel_results.new_step_size)

      # Step the inner kernel.
      new_state, new_inner_results = self.inner_kernel.one_step(
          current_state, inner_results)

      # Get the new step size.
      log_accept_prob = self.log_accept_prob_getter_fn(new_inner_results)
      log_target_accept_prob = tf.math.log(
          previous_kernel_results.target_accept_prob)

      state_parts = tf.nest.flatten(current_state)
      step_size = self.step_size_getter_fn(new_inner_results)
      step_size_parts = tf.nest.flatten(step_size)
      log_accept_prob_rank = tf.rank(log_accept_prob)

      new_step_size_parts = []
      for step_size_part, state_part in zip(step_size_parts, state_parts):
        # Compute new step sizes for each step size part. If step size part has
        # smaller rank than the corresponding state part, then the difference is
        # averaged away in the log accept prob.
        #
        # Example:
        #
        # state_part has shape      [2, 3, 4, 5]
        # step_size_part has shape     [1, 4, 1]
        # log_accept_prob has shape [2, 3, 4]
        #
        # Since step size has 1 rank fewer than the state, we reduce away the
        # leading dimension of log_accept_prob to get a Tensor with shape [3,
        # 4]. Next, since log_accept_prob must broadcast into step_size_part on
        # the left, we reduce the dimensions where their shapes differ, to get a
        # Tensor with shape [1, 4], which now is compatible with the leading
        # dimensions of step_size_part.
        #
        # There is a subtlety here in that step_size_parts might be a length-1
        # list, which means that we'll be "structure-broadcasting" it for all
        # the state parts (see logic in, e.g., hmc.py). In this case we must
        # assume that that the lone step size provided broadcasts with the event
        # dims of each state part. This means that either step size has no
        # dimensions corresponding to chain dimensions, or all states are of the
        # same shape. For the former, we want to reduce over all chain
        # dimensions. For the later, we want to use the same logic as in the
        # non-structure-broadcasted case.
        #
        # It turns out we can compute the reduction dimensions for both cases
        # uniformly by taking the rank of any state part. This obviously works
        # in the second case (where all state ranks are the same). In the first
        # case, all state parts have the rank L + D_i + B, where L is the rank
        # of log_accept_prob, D_i is the non-shared dimensions amongst all
        # states, and B are the shared dimensions of all the states, which are
        # equal to the step size. When we subtract B, we will always get a
        # number >= L, which means we'll get the full reduction we want.
        num_reduce_dims = tf.minimum(
            log_accept_prob_rank,
            tf.rank(state_part) - tf.rank(step_size_part))
        reduced_log_accept_prob = _reduce_logmeanexp(log_accept_prob,
                                                     tf.range(num_reduce_dims))
        # reduced_log_accept_prob must broadcast into step_size_part on the
        # left, so we do an additional reduction over dimensions where their
        # shapes differ.
        reduce_indices = _get_differing_dims(reduced_log_accept_prob,
                                             step_size_part)
        reduced_log_accept_prob = _reduce_logmeanexp(
            reduced_log_accept_prob, reduce_indices, keepdims=True)

        new_step_size_part = mcmc_util.choose(
            reduced_log_accept_prob > log_target_accept_prob,
            step_size_part * (1. + previous_kernel_results.adaptation_rate),
            step_size_part / (1. + previous_kernel_results.adaptation_rate))

        new_step_size_parts.append(
            tf.where(previous_kernel_results.step < self.num_adaptation_steps,
                     new_step_size_part, step_size_part))
      new_step_size = tf.nest.pack_sequence_as(step_size, new_step_size_parts)

      return new_state, previous_kernel_results._replace(
          inner_results=new_inner_results,
          step=previous_kernel_results.step + 1,
          new_step_size=new_step_size)

  def bootstrap_results(self, init_state):
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'simple_step_size_adaptation',
                                 'bootstrap_results'),
        values=[init_state]):

      inner_results = self.inner_kernel.bootstrap_results(init_state)
      step_size = self.step_size_getter_fn(inner_results)

      return SimpleStepSizeAdaptationResults(
          inner_results=inner_results,
          step=tf.constant(0, dtype=tf.int32),
          target_accept_prob=self.parameters['target_accept_prob'],
          adaptation_rate=self.parameters['adaptation_rate'],
          new_step_size=step_size)

  def is_calibrated(self):
    return self.inner_kernel.is_calibrated()

def _maybe_validate_target_accept_prob(target_accept_prob, validate_args):
  """Validates that target_accept_prob is in (0, 1)."""
  if not validate_args:
    return target_accept_prob
  with tf.control_dependencies([
      tf.compat.v1.assert_positive(
          target_accept_prob, message='`target_accept_prob` must be > 0.'),
      tf.compat.v1.assert_less(
          target_accept_prob,
          tf.ones_like(target_accept_prob),
          message='`target_accept_prob` must be < 1.')
  ]):
    return tf.identity(target_accept_prob)