with tf.compat.v1.name_scope(name, "<STR_LIT>",<EOL>[num_steps, current_state]):<EOL><INDENT>num_steps = tf.convert_to_tensor(<EOL>value=num_steps, dtype=tf.int32, name="<STR_LIT>")<EOL>if mcmc_util.is_list_like(current_state):<EOL><INDENT>current_state = [<EOL>tf.convert_to_tensor(value=s, name="<STR_LIT>")<EOL>for s in current_state<EOL>]<EOL><DEDENT>else:<EOL><INDENT>current_state = tf.convert_to_tensor(<EOL>value=current_state, name="<STR_LIT>")<EOL><DEDENT>def _make_convex_combined_log_prob_fn(iter_):<EOL><INDENT>def _fn(*args):<EOL><INDENT>p = tf.identity(proposal_log_prob_fn(*args), name="<STR_LIT>")<EOL>t = tf.identity(target_log_prob_fn(*args), name="<STR_LIT>")<EOL>dtype = p.dtype.base_dtype<EOL>beta = tf.cast(iter_ + <NUM_LIT:1>, dtype) / tf.cast(num_steps, dtype)<EOL>return tf.identity(beta * t + (<NUM_LIT:1.> - beta) * p,<EOL>name="<STR_LIT>")<EOL><DEDENT>return _fn<EOL><DEDENT>def _loop_body(iter_, ais_weights, current_state, kernel_results):<EOL><INDENT>"""<STR_LIT>"""<EOL>x = (current_state if mcmc_util.is_list_like(current_state)<EOL>else [current_state])<EOL>proposal_log_prob = proposal_log_prob_fn(*x)<EOL>target_log_prob = target_log_prob_fn(*x)<EOL>ais_weights += ((target_log_prob - proposal_log_prob) /<EOL>tf.cast(num_steps, ais_weights.dtype))<EOL>kernel = make_kernel_fn(_make_convex_combined_log_prob_fn(iter_))<EOL>next_state, inner_results = kernel.one_step(<EOL>current_state, kernel_results.inner_results)<EOL>kernel_results = AISResults(<EOL>proposal_log_prob=proposal_log_prob,<EOL>target_log_prob=target_log_prob,<EOL>inner_results=inner_results,<EOL>)<EOL>return [iter_ + <NUM_LIT:1>, ais_weights, next_state, kernel_results]<EOL><DEDENT>def _bootstrap_results(init_state):<EOL><INDENT>"""<STR_LIT>"""<EOL>kernel = make_kernel_fn(_make_convex_combined_log_prob_fn(iter_=<NUM_LIT:0>))<EOL>inner_results = kernel.bootstrap_results(init_state)<EOL>convex_combined_log_prob = inner_results.accepted_results.target_log_prob<EOL>dtype = convex_combined_log_prob.dtype.as_numpy_dtype<EOL>shape = tf.shape(input=convex_combined_log_prob)<EOL>proposal_log_prob = tf.fill(shape, dtype(np.nan),<EOL>name="<STR_LIT>")<EOL>target_log_prob = tf.fill(shape, dtype(np.nan),<EOL>name="<STR_LIT>")<EOL>return AISResults(<EOL>proposal_log_prob=proposal_log_prob,<EOL>target_log_prob=target_log_prob,<EOL>inner_results=inner_results,<EOL>)<EOL><DEDENT>previous_kernel_results = _bootstrap_results(current_state)<EOL>inner_results = previous_kernel_results.inner_results<EOL>ais_weights = tf.zeros(<EOL>shape=tf.broadcast_dynamic_shape(<EOL>tf.shape(input=inner_results.proposed_results.target_log_prob),<EOL>tf.shape(input=inner_results.accepted_results.target_log_prob)),<EOL>dtype=inner_results.proposed_results.target_log_prob.dtype.base_dtype)<EOL>[_, ais_weights, current_state, kernel_results] = tf.while_loop(<EOL>cond=lambda iter_, *args: iter_ < num_steps,<EOL>body=_loop_body,<EOL>loop_vars=[<EOL>np.int32(<NUM_LIT:0>),  <EOL>ais_weights,<EOL>current_state,<EOL>previous_kernel_results,<EOL>],<EOL>parallel_iterations=parallel_iterations)<EOL>return [current_state, ais_weights, kernel_results]<EOL><DEDENT>

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))
    ```


def _fn(state_parts, seed):<EOL><INDENT>"""<STR_LIT>"""<EOL>with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[state_parts, scale, seed]):<EOL><INDENT>scales = scale if mcmc_util.is_list_like(scale) else [scale]<EOL>if len(scales) == <NUM_LIT:1>:<EOL><INDENT>scales *= len(state_parts)<EOL><DEDENT>if len(state_parts) != len(scales):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>seed_stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>next_state_parts = [<EOL>tf.random.normal(<EOL>mean=state_part,<EOL>stddev=scale_part,<EOL>shape=tf.shape(input=state_part),<EOL>dtype=state_part.dtype.base_dtype,<EOL>seed=seed_stream())<EOL>for scale_part, state_part in zip(scales, state_parts)<EOL>]<EOL>return next_state_parts<EOL><DEDENT><DEDENT>return _fn<EOL>

Returns a callable that adds a random normal perturbation to the input.

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

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

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


def _fn(state_parts, seed):<EOL><INDENT>"""<STR_LIT>"""<EOL>with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[state_parts, scale, seed]):<EOL><INDENT>scales = scale if mcmc_util.is_list_like(scale) else [scale]<EOL>if len(scales) == <NUM_LIT:1>:<EOL><INDENT>scales *= len(state_parts)<EOL><DEDENT>if len(state_parts) != len(scales):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>seed_stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>next_state_parts = [<EOL>tf.random.uniform(<EOL>minval=state_part - scale_part,<EOL>maxval=state_part + scale_part,<EOL>shape=tf.shape(input=state_part),<EOL>dtype=state_part.dtype.base_dtype,<EOL>seed=seed_stream())<EOL>for scale_part, state_part in zip(scales, state_parts)<EOL>]<EOL>return next_state_parts<EOL><DEDENT><DEDENT>return _fn<EOL>

Returns a callable that adds a random uniform perturbation to the input.

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

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

    Returns:
      random_walk_uniform_fn: A callable accepting a Python `list` of `Tensor`s
        representing the state parts of the `current_state` and an `int`
        representing the random seed used to generate the proposal. The callable
        returns the same-type `list` of `Tensor`s as the input and represents the
        proposal for the RWM algorithm.


if mcmc_util.is_list_like(fn_arg_list):<EOL><INDENT>fn_arg_list = list(fn_arg_list)<EOL><DEDENT>else:<EOL><INDENT>fn_arg_list = [fn_arg_list]<EOL><DEDENT>if fn_result is None:<EOL><INDENT>fn_result = fn(*fn_arg_list)<EOL><DEDENT>if not fn_result.dtype.is_floating:<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>description))<EOL><DEDENT>return fn_result<EOL>

Helper which computes `fn_result` if needed.


if new_state_fn is None:<EOL><INDENT>new_state_fn = random_walk_normal_fn()<EOL><DEDENT>self._impl = metropolis_hastings.MetropolisHastings(<EOL>inner_kernel=UncalibratedRandomWalk(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>new_state_fn=new_state_fn,<EOL>seed=seed,<EOL>name=name),<EOL>seed=seed)<EOL>

Initializes this transition kernel.

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

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

        Raises:
          ValueError: if there isn't one `scale` or a list with same length as
            `current_state`.


return self._impl.inner_kernel.parameters<EOL>

Return `dict` of ``__init__`` arguments and their values.


def one_step(self, current_state, previous_kernel_results):

return self._impl.one_step(current_state, previous_kernel_results)<EOL>

Runs one iteration of Random Walk Metropolis with normal proposal.

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

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

        Raises:
          ValueError: if there isn't one `scale` or a list with same length as
            `current_state`.


return self._impl.bootstrap_results(init_state)<EOL>

Creates initial `previous_kernel_results` using a supplied `state`.


return self._parameters<EOL>


def _choose_random_direction(current_state_parts, batch_rank, seed=None):

seed_gen = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>rnd_direction_parts = [<EOL>tf.random.normal(<EOL>tf.shape(input=current_state_part), dtype=tf.float32, seed=seed_gen())<EOL>for current_state_part in current_state_parts<EOL>]<EOL>sum_squares = sum(<EOL>tf.reduce_sum(<EOL>input_tensor=rnd_direction**<NUM_LIT>,<EOL>axis=tf.range(batch_rank, tf.rank(rnd_direction)),<EOL>keepdims=True) for rnd_direction in rnd_direction_parts)<EOL>rnd_direction_parts = [rnd_direction / tf.sqrt(sum_squares)<EOL>for rnd_direction in rnd_direction_parts]<EOL>return rnd_direction_parts<EOL>

Chooses a random direction in the event space.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [<EOL>current_state_parts, step_sizes, max_doublings, current_target_log_prob,<EOL>batch_rank<EOL>]):<EOL><INDENT>direction = _choose_random_direction(current_state_parts,<EOL>batch_rank=batch_rank,<EOL>seed=seed)<EOL>reduce_axes = [tf.range(batch_rank, tf.rank(dirn_part))<EOL>for dirn_part in direction]<EOL>components = [<EOL>tf.reduce_sum(<EOL>input_tensor=(dirn_part / step_size)**<NUM_LIT:2>, axis=reduce_axes[i])<EOL>for i, (step_size, dirn_part) in enumerate(zip(step_sizes, direction))<EOL>]<EOL>step_size = tf.math.rsqrt(tf.add_n(components))<EOL>def _get_rank(x):<EOL><INDENT>return (len(x.shape.as_list()) if x.shape.dims is not None<EOL>else tf.rank(x))<EOL><DEDENT>state_part_ranks = [_get_rank(part) for part in current_state_parts]<EOL>def _step_along_direction(alpha):<EOL><INDENT>"""<STR_LIT>"""<EOL>padded_alphas = [_right_pad(alpha, final_rank=part_rank)<EOL>for part_rank in state_part_ranks]<EOL>state_parts = [state_part + padded_alpha<EOL>* direction_part for state_part, direction_part,<EOL>padded_alpha in<EOL>zip(current_state_parts, direction, padded_alphas)]<EOL>return state_parts<EOL><DEDENT>def projected_target_log_prob_fn(alpha):<EOL><INDENT>"""<STR_LIT>"""<EOL>return target_log_prob_fn(*_step_along_direction(alpha))<EOL><DEDENT>alpha_init = tf.zeros_like(current_target_log_prob,<EOL>dtype=current_state_parts[<NUM_LIT:0>].dtype.base_dtype)<EOL>[<EOL>next_alpha,<EOL>next_target_log_prob,<EOL>bounds_satisfied,<EOL>upper_bounds,<EOL>lower_bounds<EOL>] = ssu.slice_sampler_one_dim(projected_target_log_prob_fn,<EOL>x_initial=alpha_init,<EOL>max_doublings=max_doublings,<EOL>step_size=step_size, seed=seed)<EOL>return [<EOL>_step_along_direction(next_alpha),<EOL>next_target_log_prob,<EOL>bounds_satisfied,<EOL>direction,<EOL>upper_bounds,<EOL>lower_bounds<EOL>]<EOL><DEDENT>

Applies a single iteration of slice sampling update.

    Applies hit and run style slice sampling. Chooses a uniform random direction
    on the unit sphere in the event space. Applies the one dimensional slice
    sampling update along that direction.

    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `*current_state_parts` and returns its (possibly unnormalized) log-density
        under the target distribution.
      current_state_parts: Python `list` of `Tensor`s representing the current
        state(s) of the Markov chain(s). The first `independent_chain_ndims` of
        the `Tensor`(s) index different chains.
      step_sizes: Python `list` of `Tensor`s. Provides a measure of the width
        of the density. Used to find the slice bounds. Must broadcast with the
        shape of `current_state_parts`.
      max_doublings: Integer number of doublings to allow while locating the slice
        boundaries.
      current_target_log_prob: `Tensor` representing the value of
        `target_log_prob_fn(*current_state_parts)`. The only reason to specify
        this argument is to reduce TF graph size.
      batch_rank: Integer. The number of axes in the state that correspond to
        independent batches.
      seed: Python integer to seed random number generators.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'find_slice_bounds').

    Returns:
      proposed_state_parts: Tensor or Python list of `Tensor`s representing the
        state(s) of the Markov chain(s) at each result step. Has same shape as
        input `current_state_parts`.
      proposed_target_log_prob: `Tensor` representing the value of
        `target_log_prob_fn` at `next_state`.
      bounds_satisfied: Boolean `Tensor` of the same shape as the log density.
        True indicates whether the an interval containing the slice for that
        batch was found successfully.
      direction: `Tensor` or Python list of `Tensors`s representing the direction
        along which the slice was sampled. Has the same shape and dtype(s) as
        `current_state_parts`.
      upper_bounds: `Tensor` of batch shape and the dtype of the input state. The
        upper bounds of the slices along the sampling direction.
      lower_bounds: `Tensor` of batch shape and the dtype of the input state. The
        lower bounds of the slices along the sampling direction.


fn_arg_list = (list(fn_arg_list) if mcmc_util.is_list_like(fn_arg_list)<EOL>else [fn_arg_list])<EOL>if fn_result is None:<EOL><INDENT>fn_result = fn(*fn_arg_list)<EOL><DEDENT>if not fn_result.dtype.is_floating:<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>description))<EOL><DEDENT>return fn_result<EOL>


padded_shape = tf.concat(<EOL>[tf.shape(input=x),<EOL>tf.ones(final_rank - tf.rank(x), dtype=tf.int32)],<EOL>axis=<NUM_LIT:0>)<EOL>static_padded_shape = None<EOL>if x.shape.is_fully_defined() and isinstance(final_rank, int):<EOL><INDENT>static_padded_shape = x.shape.as_list()<EOL>extra_dims = final_rank - len(static_padded_shape)<EOL>static_padded_shape.extend([<NUM_LIT:1>] * extra_dims)<EOL><DEDENT>padded_x = tf.reshape(x, static_padded_shape or padded_shape)<EOL>return padded_x<EOL>

Pads the shape of x to the right to be of rank final_rank.

    Expands the dims of `x` to the right such that its rank is equal to
    final_rank. For example, if `x` is of shape [1, 5, 7, 2] and `final_rank` is
    7, we return padded_x, which is of shape [1, 5, 7, 2, 1, 1, 1].

    Args:
      x: The tensor whose shape is to be padded.
      final_rank: Scalar int32 `Tensor` or Python `int`. The desired rank of x.

    Returns:
      padded_x: A tensor of rank final_rank.


def _prepare_args(target_log_prob_fn, state, step_size,<EOL>target_log_prob=None, maybe_expand=False,<EOL>description='<STR_LIT>'):

state_parts = list(state) if mcmc_util.is_list_like(state) else [state]<EOL>state_parts = [<EOL>tf.convert_to_tensor(value=s, name='<STR_LIT>') for s in state_parts<EOL>]<EOL>target_log_prob = _maybe_call_fn(<EOL>target_log_prob_fn,<EOL>state_parts,<EOL>target_log_prob,<EOL>description)<EOL>step_sizes = (list(step_size) if mcmc_util.is_list_like(step_size)<EOL>else [step_size])<EOL>step_sizes = [<EOL>tf.convert_to_tensor(<EOL>value=s, name='<STR_LIT>', dtype=target_log_prob.dtype)<EOL>for s in step_sizes<EOL>]<EOL>if len(step_sizes) == <NUM_LIT:1>:<EOL><INDENT>step_sizes *= len(state_parts)<EOL><DEDENT>if len(state_parts) != len(step_sizes):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>def maybe_flatten(x):<EOL><INDENT>return x if maybe_expand or mcmc_util.is_list_like(state) else x[<NUM_LIT:0>]<EOL><DEDENT>return [<EOL>maybe_flatten(state_parts),<EOL>maybe_flatten(step_sizes),<EOL>target_log_prob<EOL>]<EOL>

Processes input args to meet list-like assumptions.


self._seed_stream = distributions.SeedStream(<EOL>seed, salt='<STR_LIT>')<EOL>self._parameters = dict(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>step_size=step_size,<EOL>max_doublings=max_doublings,<EOL>seed=seed,<EOL>name=name)<EOL>

Initializes this transition kernel.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it is a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          step_size: Scalar or `tf.Tensor` with same dtype as and shape compatible
            with `x_initial`. The size of the initial interval.
          max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
          doublings to consider.
          seed: Python integer to seed the random number generator.
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., 'slice_sampler_kernel').

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


return self._parameters<EOL>

Returns `dict` of ``__init__`` arguments and their values.


Runs one iteration of Slice Sampler.

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

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

        Raises:
          ValueError: if there isn't one `step_size` or a list with same length as
            `current_state`.
          TypeError: if `not target_log_prob.dtype.is_floating`.


def tridiag(d, diag_value, offdiag_value):

diag_mat = tf.eye(d) * (diag_value - offdiag_value)<EOL>three_bands = tf.linalg.band_part(tf.fill([d, d], offdiag_value), <NUM_LIT:1>, <NUM_LIT:1>)<EOL>return diag_mat + three_bands<EOL>

d x d matrix with given value on diag, and one super/sub diag.


<EOL>logu = np.float128(logu)<EOL>n = logu.shape[<NUM_LIT:0>]<EOL>u = np.exp(logu)<EOL>loogeoavg_u = []  <EOL>for j in range(n):<EOL><INDENT>loogeoavg_u.append(np.exp(np.mean(<EOL>[logu[i, ...] for i in range(n) if i != j],<EOL>axis=<NUM_LIT:0>)))<EOL><DEDENT>loogeoavg_u = np.array(loogeoavg_u)<EOL>loosum_u = []  <EOL>for j in range(n):<EOL><INDENT>loosum_u.append(np.sum(<EOL>[u[i, ...] for i in range(n) if i != j],<EOL>axis=<NUM_LIT:0>))<EOL><DEDENT>loosum_u = np.array(loosum_u)<EOL>log_sooavg_u = np.log(loosum_u + loogeoavg_u) - np.log(n)<EOL>log_avg_u = np.log(np.mean(u, axis=<NUM_LIT:0>))<EOL>return log_avg_u, log_sooavg_u<EOL>

Numpy implementation of `csiszar_vimco_helper`.


def _csiszar_vimco_helper_grad(self, logu, delta):

<EOL>np_log_avg_u1, np_log_sooavg_u1 = self._csiszar_vimco_helper(<EOL>logu[..., None] + np.diag([delta]*len(logu)))<EOL>np_log_avg_u, np_log_sooavg_u = self._csiszar_vimco_helper(<EOL>logu[..., None])<EOL>return [<EOL>(np_log_avg_u1 - np_log_avg_u) / delta,<EOL>np.sum(np_log_sooavg_u1 - np_log_sooavg_u, axis=<NUM_LIT:0>) / delta,<EOL>]<EOL>

Finite difference approximation of `grad(csiszar_vimco_helper, logu)`.


def amari_alpha(logu, alpha=<NUM_LIT:1.>, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>if alpha is None or tf.is_tensor(alpha):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if (self_normalized is None or tf.is_tensor(self_normalized)):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>if alpha == <NUM_LIT:0.>:<EOL><INDENT>f = -logu<EOL><DEDENT>elif alpha == <NUM_LIT:1.>:<EOL><INDENT>f = tf.exp(logu) * logu<EOL><DEDENT>else:<EOL><INDENT>f = tf.math.expm1(alpha * logu) / (alpha * (alpha - <NUM_LIT:1.>))<EOL><DEDENT>if not self_normalized:<EOL><INDENT>return f<EOL><DEDENT>if alpha == <NUM_LIT:0.>:<EOL><INDENT>return f + tf.math.expm1(logu)<EOL><DEDENT>elif alpha == <NUM_LIT:1.>:<EOL><INDENT>return f - tf.math.expm1(logu)<EOL><DEDENT>else:<EOL><INDENT>return f - tf.math.expm1(logu) / (alpha - <NUM_LIT:1.>)<EOL><DEDENT><DEDENT>

The Amari-alpha Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True`, the Amari-alpha Csiszar-function is:

    ```none
    f(u) = { -log(u) + (u - 1),     alpha = 0
           { u log(u) - (u - 1),    alpha = 1
           { [(u**alpha - 1) - alpha (u - 1)] / (alpha (alpha - 1)),    otherwise
    ```

    When `self_normalized = False` the `(u - 1)` terms are omitted.

    Warning: when `alpha != 0` and/or `self_normalized = True` this function makes
    non-log-space calculations and may therefore be numerically unstable for
    `|logu| >> 0`.

    For more information, see:
      A. Cichocki and S. Amari. "Families of Alpha-Beta-and GammaDivergences:
      Flexible and Robust Measures of Similarities." Entropy, vol. 12, no. 6, pp.
      1532-1568, 2010.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      alpha: `float`-like Python scalar. (See Mathematical Details for meaning.)
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      amari_alpha_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.

    Raises:
      TypeError: if `alpha` is `None` or a `Tensor`.
      TypeError: if `self_normalized` is `None` or a `Tensor`.


def kl_reverse(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>return amari_alpha(logu, alpha=<NUM_LIT:0.>, self_normalized=self_normalized)<EOL><DEDENT>

The reverse Kullback-Leibler Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True`, the KL-reverse Csiszar-function is:

    ```none
    f(u) = -log(u) + (u - 1)
    ```

    When `self_normalized = False` the `(u - 1)` term is omitted.

    Observe that as an f-Divergence, this Csiszar-function implies:

    ```none
    D_f[p, q] = KL[q, p]
    ```

    The KL is "reverse" because in maximum likelihood we think of minimizing `q`
    as in `KL[p, q]`.

    Warning: when self_normalized = True` this function makes non-log-space
    calculations and may therefore be numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      kl_reverse_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
        `u = exp(logu)`.

    Raises:
      TypeError: if `self_normalized` is `None` or a `Tensor`.


def kl_forward(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>return amari_alpha(logu, alpha=<NUM_LIT:1.>, self_normalized=self_normalized)<EOL><DEDENT>

The forward Kullback-Leibler Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True`, the KL-forward Csiszar-function is:

    ```none
    f(u) = u log(u) - (u - 1)
    ```

    When `self_normalized = False` the `(u - 1)` term is omitted.

    Observe that as an f-Divergence, this Csiszar-function implies:

    ```none
    D_f[p, q] = KL[p, q]
    ```

    The KL is "forward" because in maximum likelihood we think of minimizing `q`
    as in `KL[p, q]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      kl_forward_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
        `u = exp(logu)`.

    Raises:
      TypeError: if `self_normalized` is `None` or a `Tensor`.


def jensen_shannon(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>npdt = logu.dtype.as_numpy_dtype<EOL>y = tf.nn.softplus(logu)<EOL>if self_normalized:<EOL><INDENT>y -= np.log(<NUM_LIT:2>).astype(npdt)<EOL><DEDENT>return tf.exp(logu) * logu - (<NUM_LIT:1.> + tf.exp(logu)) * y<EOL><DEDENT>

The Jensen-Shannon Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True`, the Jensen-Shannon Csiszar-function is:

    ```none
    f(u) = u log(u) - (1 + u) log(1 + u) + (u + 1) log(2)
    ```

    When `self_normalized = False` the `(u + 1) log(2)` term is omitted.

    Observe that as an f-Divergence, this Csiszar-function implies:

    ```none
    D_f[p, q] = KL[p, m] + KL[q, m]
    m(x) = 0.5 p(x) + 0.5 q(x)
    ```

    In a sense, this divergence is the "reverse" of the Arithmetic-Geometric
    f-Divergence.

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    For more information, see:
      Lin, J. "Divergence measures based on the Shannon entropy." IEEE Trans.
      Inf. Th., 37, 145-151, 1991.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      jensen_shannon_of_u: `float`-like `Tensor` of the Csiszar-function
        evaluated at `u = exp(logu)`.


def arithmetic_geometric(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>y = tf.nn.softplus(logu) - <NUM_LIT:0.5> * logu<EOL>if self_normalized:<EOL><INDENT>y -= np.log(<NUM_LIT>).astype(logu.dtype.as_numpy_dtype)<EOL><DEDENT>return (<NUM_LIT:1.> + tf.exp(logu)) * y<EOL><DEDENT>

The Arithmetic-Geometric Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True` the Arithmetic-Geometric Csiszar-function is:

    ```none
    f(u) = (1 + u) log( (1 + u) / sqrt(u) ) - (1 + u) log(2)
    ```

    When `self_normalized = False` the `(1 + u) log(2)` term is omitted.

    Observe that as an f-Divergence, this Csiszar-function implies:

    ```none
    D_f[p, q] = KL[m, p] + KL[m, q]
    m(x) = 0.5 p(x) + 0.5 q(x)
    ```

    In a sense, this divergence is the "reverse" of the Jensen-Shannon
    f-Divergence.

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      arithmetic_geometric_of_u: `float`-like `Tensor` of the
        Csiszar-function evaluated at `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return <NUM_LIT:0.5> * tf.abs(tf.math.expm1(logu))<EOL><DEDENT>

The Total Variation Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Total-Variation Csiszar-function is:

    ```none
    f(u) = 0.5 |u - 1|
    ```

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      total_variation_of_u: `float`-like `Tensor` of the Csiszar-function
        evaluated at `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return tf.square(tf.math.expm1(logu))<EOL><DEDENT>

The Pearson Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Pearson Csiszar-function is:

    ```none
    f(u) = (u - 1)**2
    ```

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      pearson_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
        `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return pearson(<NUM_LIT:0.5> * logu)<EOL><DEDENT>

The Squared-Hellinger Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Squared-Hellinger Csiszar-function is:

    ```none
    f(u) = (sqrt(u) - 1)**2
    ```

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      squared_hellinger_of_u: `float`-like `Tensor` of the Csiszar-function
        evaluated at `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return pearson(logu) / (<NUM_LIT:1.> + tf.exp(logu))<EOL><DEDENT>

The Triangular Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Triangular Csiszar-function is:

    ```none
    f(u) = (u - 1)**2 / (1 + u)
    ```

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      triangular_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


def t_power(logu, t, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu, t]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>t = tf.convert_to_tensor(value=t, dtype=logu.dtype.base_dtype, name="<STR_LIT:t>")<EOL>fu = tf.math.expm1(t * logu)<EOL>if self_normalized:<EOL><INDENT>fu -= t * tf.math.expm1(logu)<EOL><DEDENT>fu *= tf.where(tf.logical_and(<NUM_LIT:0.> < t, t < <NUM_LIT:1.>),<EOL>-tf.ones_like(t),<EOL>tf.ones_like(t))<EOL>return fu<EOL><DEDENT>

The T-Power Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True` the T-Power Csiszar-function is:

    ```none
    f(u) = s [ u**t - 1 - t(u - 1) ]
    s = { -1   0 < t < 1
        { +1   otherwise
    ```

    When `self_normalized = False` the `- t(u - 1)` term is omitted.

    This is similar to the `amari_alpha` Csiszar-function, with the associated
    divergence being the same up to factors depending only on `t`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      t:  `Tensor` of same `dtype` as `logu` and broadcastable shape.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      t_power_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return tf.math.expm1(tf.abs(logu))<EOL><DEDENT>

The log1p-abs Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Log1p-Abs Csiszar-function is:

    ```none
    f(u) = u**(sign(u-1)) - 1
    ```

    This function is so-named because it was invented from the following recipe.
    Choose a convex function g such that g(0)=0 and solve for f:

    ```none
    log(1 + f(u)) = g(log(u)).
      <=>
    f(u) = exp(g(log(u))) - 1
    ```

    That is, the graph is identically `g` when y-axis is `log1p`-domain and x-axis
    is `log`-domain.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      log1p_abs_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return <NUM_LIT:0.5> * tf.math.expm1(logu) * logu<EOL><DEDENT>

The Jeffreys Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Jeffreys Csiszar-function is:

    ```none
    f(u) = 0.5 ( u log(u) - log(u) )
         = 0.5 kl_forward + 0.5 kl_reverse
         = symmetrized_csiszar_function(kl_reverse)
         = symmetrized_csiszar_function(kl_forward)
    ```

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      jeffreys_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return tf.math.expm1(<NUM_LIT> * logu)<EOL><DEDENT>

The chi-Square Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Chi-square Csiszar-function is:

    ```none
    f(u) = u**2 - 1
    ```

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


def modified_gan(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>y = tf.nn.softplus(logu) - logu<EOL>if self_normalized:<EOL><INDENT>y += <NUM_LIT:0.5> * tf.math.expm1(logu)<EOL><DEDENT>return y<EOL><DEDENT>

The Modified-GAN Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True` the modified-GAN (Generative/Adversarial
    Network) Csiszar-function is:

    ```none
    f(u) = log(1 + u) - log(u) + 0.5 (u - 1)
    ```

    When `self_normalized = False` the `0.5 (u - 1)` is omitted.

    The unmodified GAN Csiszar-function is identical to Jensen-Shannon (with
    `self_normalized = False`).

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


def dual_csiszar_function(logu, csiszar_function, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>return tf.exp(logu) * csiszar_function(-logu)<EOL><DEDENT>

Calculates the dual Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Csiszar-dual is defined as:

    ```none
    f^*(u) = u f(1 / u)
    ```

    where `f` is some other Csiszar-function.

    For example, the dual of `kl_reverse` is `kl_forward`, i.e.,

    ```none
    f(u) = -log(u)
    f^*(u) = u f(1 / u) = -u log(1 / u) = u log(u)
    ```

    The dual of the dual is the original function:

    ```none
    f^**(u) = {u f(1/u)}^*(u) = u (1/u) f(1/(1/u)) = f(u)
    ```

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      csiszar_function: Python `callable` representing a Csiszar-function over
        log-domain.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      dual_f_of_u: `float`-like `Tensor` of the result of calculating the dual of
        `f` at `u = exp(logu)`.


def symmetrized_csiszar_function(logu, csiszar_function, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return <NUM_LIT:0.5> * (csiszar_function(logu)<EOL>+ dual_csiszar_function(logu, csiszar_function))<EOL><DEDENT>

Symmetrizes a Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The symmetrized Csiszar-function is defined as:

    ```none
    f_g(u) = 0.5 g(u) + 0.5 u g (1 / u)
    ```

    where `g` is some other Csiszar-function.

    We say the function is "symmetrized" because:

    ```none
    D_{f_g}[p, q] = D_{f_g}[q, p]
    ```

    for all `p << >> q` (i.e., `support(p) = support(q)`).

    There exists alternatives for symmetrizing a Csiszar-function. For example,

    ```none
    f_g(u) = max(f(u), f^*(u)),
    ```

    where `f^*` is the dual Csiszar-function, also implies a symmetric
    f-Divergence.

    Example:

    When either of the following functions are symmetrized, we obtain the
    Jensen-Shannon Csiszar-function, i.e.,

    ```none
    g(u) = -log(u) - (1 + u) log((1 + u) / 2) + u - 1
    h(u) = log(4) + 2 u log(u / (1 + u))
    ```

    implies,

    ```none
    f_g(u) = f_h(u) = u log(u) - (1 + u) log((1 + u) / 2)
           = jensen_shannon(log(u)).
    ```

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      csiszar_function: Python `callable` representing a Csiszar-function over
        log-domain.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      symmetrized_g_of_u: `float`-like `Tensor` of the result of applying the
        symmetrization of `g` evaluated at `u = exp(logu)`.


reparameterization_types = tf.nest.flatten(q.reparameterization_type)<EOL>with tf.compat.v1.name_scope(name, "<STR_LIT>",<EOL>[num_draws]):<EOL><INDENT>if use_reparametrization is None:<EOL><INDENT>use_reparametrization = all(<EOL>reparameterization_type == tfd.FULLY_REPARAMETERIZED<EOL>for reparameterization_type in reparameterization_types)<EOL><DEDENT>elif (use_reparametrization and<EOL>any(reparameterization_type != tfd.FULLY_REPARAMETERIZED<EOL>for reparameterization_type in reparameterization_types)):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>if not callable(p_log_prob):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>return monte_carlo.expectation(<EOL>f=lambda q_samples: f(p_log_prob(q_samples) - q.log_prob(q_samples)),<EOL>samples=q.sample(num_draws, seed=seed),<EOL>log_prob=q.log_prob,  <EOL>use_reparametrization=use_reparametrization)<EOL><DEDENT>

Monte-Carlo approximation of the Csiszar f-Divergence.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Csiszar f-Divergence for Csiszar-function f is given by:

    ```none
    D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ]
                    ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ),
                               where x_j ~iid q(X)
    ```

    Tricks: Reparameterization and Score-Gradient

    When q is "reparameterized", i.e., a diffeomorphic transformation of a
    parameterless distribution (e.g.,
    `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and
    expectation, i.e.,
    `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}`
    and `s_i = f(x_i), x_i ~iid q(X)`.

    However, if q is not reparameterized, TensorFlow's gradient will be incorrect
    since the chain-rule stops at samples of unreparameterized distributions. In
    this circumstance using the Score-Gradient trick results in an unbiased
    gradient, i.e.,

    ```none
    grad[ E_q[f(X)] ]
    = grad[ int dx q(x) f(x) ]
    = int dx grad[ q(x) f(x) ]
    = int dx [ q'(x) f(x) + q(x) f'(x) ]
    = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ]
    = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ]
    = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ]
    ```

    Unless `q.reparameterization_type != tfd.FULLY_REPARAMETERIZED` it is
    usually preferable to set `use_reparametrization = True`.

    Example Application:

    The Csiszar f-Divergence is a useful framework for variational inference.
    I.e., observe that,

    ```none
    f(p(x)) =  f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] )
            <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ]
            := D_f[p(x, Z), q(Z | x)]
    ```

    The inequality follows from the fact that the "perspective" of `f`, i.e.,
    `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and
    `t` is a real. Since the above framework includes the popular Evidence Lower
    BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework
    "Evidence Divergence Bound Optimization" (EDBO).

    Args:
      f: Python `callable` representing a Csiszar-function in log-space, i.e.,
        takes `p_log_prob(q_samples) - q.log_prob(q_samples)`.
      p_log_prob: Python `callable` taking (a batch of) samples from `q` and
        returning the natural-log of the probability under distribution `p`.
        (In variational inference `p` is the joint distribution.)
      q: `tf.Distribution`-like instance; must implement:
        `reparameterization_type`, `sample(n, seed)`, and `log_prob(x)`.
        (In variational inference `q` is the approximate posterior distribution.)
      num_draws: Integer scalar number of draws used to approximate the
        f-Divergence expectation.
      use_reparametrization: Python `bool`. When `None` (the default),
        automatically set to:
        `q.reparameterization_type == tfd.FULLY_REPARAMETERIZED`.
        When `True` uses the standard Monte-Carlo average. When `False` uses the
        score-gradient trick. (See above for details.)  When `False`, consider
        using `csiszar_vimco`.
      seed: Python `int` seed for `q.sample`.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      monte_carlo_csiszar_f_divergence: `float`-like `Tensor` Monte Carlo
        approximation of the Csiszar f-Divergence.

    Raises:
      ValueError: if `q` is not a reparameterized distribution and
        `use_reparametrization = True`. A distribution `q` is said to be
        "reparameterized" when its samples are generated by transforming the
        samples of another distribution which does not depend on the
        parameterization of `q`. This property ensures the gradient (with respect
        to parameters) is valid.
      TypeError: if `p_log_prob` is not a Python `callable`.


with tf.compat.v1.name_scope(name, "<STR_LIT>",<EOL>[num_draws, num_batch_draws]):<EOL><INDENT>if num_draws < <NUM_LIT:2>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>stop = tf.stop_gradient  <EOL>x = stop(q.sample(sample_shape=[num_draws, num_batch_draws],<EOL>seed=seed))<EOL>logqx = q.log_prob(x)<EOL>logu = p_log_prob(x) - logqx<EOL>f_log_avg_u, f_log_sooavg_u = [f(r) for r in csiszar_vimco_helper(logu)]<EOL>dotprod = tf.reduce_sum(<EOL>input_tensor=logqx * stop(f_log_avg_u - f_log_sooavg_u),<EOL>axis=<NUM_LIT:0>)  <EOL>f_log_avg_u += dotprod - stop(dotprod)  <EOL>return tf.reduce_mean(input_tensor=f_log_avg_u, axis=<NUM_LIT:0>)<EOL><DEDENT>

Use VIMCO to lower the variance of gradient[csiszar_function(Avg(logu))].

    This function generalizes VIMCO [(Mnih and Rezende, 2016)][1] to Csiszar
    f-Divergences.

    Note: if `q.reparameterization_type = tfd.FULLY_REPARAMETERIZED`,
    consider using `monte_carlo_csiszar_f_divergence`.

    The VIMCO loss is:

    ```none
    vimco = f(Avg{logu[i] : i=0,...,m-1})
    where,
      logu[i] = log( p(x, h[i]) / q(h[i] | x) )
      h[i] iid~ q(H | x)
    ```

    Interestingly, the VIMCO gradient is not the naive gradient of `vimco`.
    Rather, it is characterized by:

    ```none
    grad[vimco] - variance_reducing_term
    where,
      variance_reducing_term = Sum{ grad[log q(h[i] | x)] *
                                      (vimco - f(log Avg{h[j;i] : j=0,...,m-1}))
                                   : i=0, ..., m-1 }
      h[j;i] = { u[j]                             j!=i
               { GeometricAverage{ u[k] : k!=i}   j==i
    ```

    (We omitted `stop_gradient` for brevity. See implementation for more details.)

    The `Avg{h[j;i] : j}` term is a kind of "swap-out average" where the `i`-th
    element has been replaced by the leave-`i`-out Geometric-average.

    This implementation prefers numerical precision over efficiency, i.e.,
    `O(num_draws * num_batch_draws * prod(batch_shape) * prod(event_shape))`.
    (The constant may be fairly large, perhaps around 12.)

    Args:
      f: Python `callable` representing a Csiszar-function in log-space.
      p_log_prob: Python `callable` representing the natural-log of the
        probability under distribution `p`. (In variational inference `p` is the
        joint distribution.)
      q: `tf.Distribution`-like instance; must implement: `sample(n, seed)`, and
        `log_prob(x)`. (In variational inference `q` is the approximate posterior
        distribution.)
      num_draws: Integer scalar number of draws used to approximate the
        f-Divergence expectation.
      num_batch_draws: Integer scalar number of draws used to approximate the
        f-Divergence expectation.
      seed: Python `int` seed for `q.sample`.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      vimco: The Csiszar f-Divergence generalized VIMCO objective.

    Raises:
      ValueError: if `num_draws < 2`.

    #### References

    [1]: Andriy Mnih and Danilo Rezende. Variational Inference for Monte Carlo
         objectives. In _International Conference on Machine Learning_, 2016.
         https://arxiv.org/abs/1602.06725


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>n = tf.compat.dimension_value(logu.shape.with_rank_at_least(<NUM_LIT:1>)[<NUM_LIT:0>])<EOL>if n is None:<EOL><INDENT>n = tf.shape(input=logu)[<NUM_LIT:0>]<EOL>log_n = tf.math.log(tf.cast(n, dtype=logu.dtype))<EOL>nm1 = tf.cast(n - <NUM_LIT:1>, dtype=logu.dtype)<EOL><DEDENT>else:<EOL><INDENT>log_n = np.log(n).astype(logu.dtype.as_numpy_dtype)<EOL>nm1 = np.asarray(n - <NUM_LIT:1>, dtype=logu.dtype.as_numpy_dtype)<EOL><DEDENT>log_max_u = tf.reduce_max(input_tensor=logu, axis=<NUM_LIT:0>)<EOL>log_sum_u_minus_log_max_u = tf.reduce_logsumexp(<EOL>input_tensor=logu - log_max_u, axis=<NUM_LIT:0>)<EOL>d = log_sum_u_minus_log_max_u + (log_max_u - logu)<EOL>d_ok = tf.not_equal(d, <NUM_LIT:0.>)<EOL>safe_d = tf.where(d_ok, d, tf.ones_like(d))<EOL>d_ok_result = logu + tfd.softplus_inverse(safe_d)<EOL>inf = np.array(np.inf, dtype=logu.dtype.as_numpy_dtype)<EOL>is_positive_and_largest = tf.logical_and(<EOL>logu > <NUM_LIT:0.>,<EOL>tf.equal(logu, log_max_u[tf.newaxis, ...]))<EOL>log_lomsum_u = tf.reduce_logsumexp(<EOL>input_tensor=tf.where(is_positive_and_largest,<EOL>tf.fill(tf.shape(input=logu), -inf), logu),<EOL>axis=<NUM_LIT:0>,<EOL>keepdims=True)<EOL>log_lomsum_u = tf.tile(<EOL>log_lomsum_u,<EOL>multiples=<NUM_LIT:1> + tf.pad(tensor=[n - <NUM_LIT:1>], paddings=[[<NUM_LIT:0>, tf.rank(logu) - <NUM_LIT:1>]]))<EOL>d_not_ok_result = tf.where(is_positive_and_largest, log_lomsum_u,<EOL>tf.fill(tf.shape(input=d), -inf))<EOL>log_loosum_u = tf.where(d_ok, d_ok_result, d_not_ok_result)<EOL>looavg_logu = (tf.reduce_sum(input_tensor=logu, axis=<NUM_LIT:0>) - logu) / nm1<EOL>log_soosum_u = tf.reduce_logsumexp(<EOL>input_tensor=tf.stack([log_loosum_u, looavg_logu]), axis=<NUM_LIT:0>)<EOL>log_avg_u = log_sum_u_minus_log_max_u + log_max_u - log_n<EOL>log_sooavg_u = log_soosum_u - log_n<EOL>log_avg_u.set_shape(logu.shape.with_rank_at_least(<NUM_LIT:1>)[<NUM_LIT:1>:])<EOL>log_sooavg_u.set_shape(logu.shape)<EOL>return log_avg_u, log_sooavg_u<EOL><DEDENT>

Helper to `csiszar_vimco`; computes `log_avg_u`, `log_sooavg_u`.

    `axis = 0` of `logu` is presumed to correspond to iid samples from `q`, i.e.,

    ```none
    logu[j] = log(u[j])
    u[j] = p(x, h[j]) / q(h[j] | x)
    h[j] iid~ q(H | x)
    ```

    Args:
      logu: Floating-type `Tensor` representing `log(p(x, h) / q(h | x))`.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      log_avg_u: `logu.dtype` `Tensor` corresponding to the natural-log of the
        average of `u`. The sum of the gradient of `log_avg_u` is `1`.
      log_sooavg_u: `logu.dtype` `Tensor` characterized by the natural-log of the
        average of `u`` except that the average swaps-out `u[i]` for the
        leave-`i`-out Geometric-average. The mean of the gradient of
        `log_sooavg_u` is `1`. Mathematically `log_sooavg_u` is,
        ```none
        log_sooavg_u[i] = log(Avg{h[j ; i] : j=0, ..., m-1})
        h[j ; i] = { u[j]                              j!=i
                   { GeometricAverage{u[k] : k != i}   j==i
        ```


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[population, population_values, differential_weight, crossover_prob]):<EOL><INDENT>population, _ = _ensure_list(population)<EOL>if population_values is None:<EOL><INDENT>population_values = objective_function(*population)<EOL><DEDENT>population_size = tf.shape(input=population[<NUM_LIT:0>])[<NUM_LIT:0>]<EOL>seed_stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>mixing_indices = _get_mixing_indices(population_size, seed=seed_stream())<EOL>mutants = _get_mutants(population,<EOL>population_size,<EOL>mixing_indices,<EOL>differential_weight)<EOL>candidates = _binary_crossover(population,<EOL>population_size,<EOL>mutants,<EOL>crossover_prob,<EOL>seed=seed_stream())<EOL>candidate_values = objective_function(*candidates)<EOL>if population_values is None:<EOL><INDENT>population_values = objective_function(*population)<EOL><DEDENT>infinity = tf.zeros_like(population_values) + np.inf<EOL>population_values = tf.where(<EOL>tf.math.is_nan(population_values), x=infinity, y=population_values)<EOL>to_replace = candidate_values < population_values<EOL>next_population = [<EOL>tf.where(to_replace, x=candidates_part, y=population_part)<EOL>for candidates_part, population_part in zip(candidates, population)<EOL>]<EOL>next_values = tf.where(to_replace, x=candidate_values, y=population_values)<EOL><DEDENT>return next_population, next_values<EOL>

Performs one step of the differential evolution algorithm.

    Args:
      objective_function:  A Python callable that accepts a batch of possible
        solutions and returns the values of the objective function at those
        arguments as a rank 1 real `Tensor`. This specifies the function to be
        minimized. The input to this callable may be either a single `Tensor`
        or a Python `list` of `Tensor`s. The signature must match the format of
        the argument `population`. (i.e. objective_function(*population) must
        return the value of the function to be minimized).
      population:  `Tensor` or Python `list` of `Tensor`s representing the
        current population vectors. Each `Tensor` must be of the same real dtype.
        The first dimension indexes individual population members while the
        rest of the dimensions are consumed by the value function. For example,
        if the population is a single `Tensor` of shape [n, m1, m2], then `n` is
        the population size and the output of `objective_function` applied to the
        population is a `Tensor` of shape [n]. If the population is a python
        list of `Tensor`s then each `Tensor` in the list should have the first
        axis of a common size, say `n` and `objective_function(*population)`
        should return a `Tensor of shape [n]. The population must have at least
        4 members for the algorithm to work correctly.
      population_values: A `Tensor` of rank 1 and real dtype. The result of
        applying `objective_function` to the `population`. If not supplied it is
        computed using the `objective_function`.
        Default value: None.
      differential_weight: Real scalar `Tensor`. Must be positive and less than
        2.0. The parameter controlling the strength of mutation.
        Default value: 0.5
      crossover_prob: Real scalar `Tensor`. Must be between 0 and 1. The
        probability of recombination per site.
        Default value: 0.9
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: None.
      name: (Optional) Python str. The name prefixed to the ops created by this
        function. If not supplied, the default name 'one_step' is
        used.
        Default value: None

    Returns:
      A sequence containing the following elements (in order):
      next_population: A `Tensor` or Python `list` of `Tensor`s of the same
        structure as the input population. The population at the next generation.
      next_population_values: A `Tensor` of same shape and dtype as input
        `population_values`. The function values for the `next_population`.


if initial_population is None and initial_position is None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if initial_population is not None and initial_position is not None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name,<EOL>default_name='<STR_LIT>',<EOL>values=[<EOL>initial_population, initial_position, population_size,<EOL>population_stddev, max_iterations, func_tolerance, position_tolerance,<EOL>differential_weight, crossover_prob<EOL>]):<EOL><INDENT>(<EOL>was_iterable,<EOL>population,<EOL>population_values,<EOL>max_iterations,<EOL>func_tolerance,<EOL>position_tolerance,<EOL>differential_weight,<EOL>crossover_prob<EOL>) = _get_initial_args(objective_function,<EOL>initial_population,<EOL>initial_position,<EOL>population_size,<EOL>population_stddev,<EOL>max_iterations,<EOL>func_tolerance,<EOL>position_tolerance,<EOL>differential_weight,<EOL>crossover_prob,<EOL>seed)<EOL>def evolve_body(loop_vars):<EOL><INDENT>"""<STR_LIT>"""<EOL>next_population, next_population_values = one_step(<EOL>objective_function,<EOL>loop_vars.population,<EOL>population_values=loop_vars.population_values,<EOL>differential_weight=differential_weight,<EOL>crossover_prob=crossover_prob,<EOL>seed=seed)<EOL>converged = _check_convergence(next_population,<EOL>next_population_values,<EOL>func_tolerance,<EOL>position_tolerance)<EOL>failed = _check_failure(next_population_values)<EOL>return [_MinimizeLoopVars(<EOL>converged=converged,<EOL>failed=failed,<EOL>num_iterations=loop_vars.num_iterations+<NUM_LIT:1>,<EOL>population=next_population,<EOL>population_values=next_population_values)]<EOL><DEDENT>def evolve_cond(loop_vars):<EOL><INDENT>should_stop = (<EOL>loop_vars.failed |<EOL>loop_vars.converged |<EOL>(max_iterations is not None and<EOL>loop_vars.num_iterations >= max_iterations))<EOL>return ~should_stop<EOL><DEDENT>initial_vars = _MinimizeLoopVars(<EOL>converged=tf.convert_to_tensor(value=False),<EOL>failed=tf.convert_to_tensor(value=False),<EOL>num_iterations=tf.convert_to_tensor(value=<NUM_LIT:0>),<EOL>population=population,<EOL>population_values=population_values)<EOL>final_state = tf.while_loop(<EOL>cond=evolve_cond, body=evolve_body, loop_vars=(initial_vars,))[<NUM_LIT:0>]<EOL>best_position, best_values = _find_best_in_population(<EOL>final_state.population,<EOL>final_state.population_values)<EOL>final_population = final_state.population<EOL>if not was_iterable:<EOL><INDENT>final_population = final_population[<NUM_LIT:0>]<EOL>best_position = best_position[<NUM_LIT:0>]<EOL><DEDENT>return DifferentialEvolutionOptimizerResults(<EOL>converged=final_state.converged,<EOL>failed=final_state.failed,<EOL>position=best_position,<EOL>objective_value=best_values,<EOL>final_population=final_population,<EOL>final_objective_values=final_state.population_values,<EOL>initial_population=population,<EOL>initial_objective_values=population_values,<EOL>num_iterations=final_state.num_iterations)<EOL><DEDENT>

Applies the Differential evolution algorithm to minimize a function.

    Differential Evolution is an evolutionary optimization algorithm which works
    on a set of candidate solutions called the population. It iteratively
    improves the population by applying genetic operators of mutation and
    recombination. The objective function `f` supplies the fitness of each
    candidate. A candidate `s_1` is considered better than `s_2` if
    `f(s_1) < f(s_2)`.

    This method allows the user to either specify an initial population or a
    single candidate solution. If a single solution is specified, a population
    of the specified size is initialized by adding independent normal noise
    to the candidate solution.

    The implementation also supports a multi-part specification of the state. For
    example, consider the objective function:

    ```python
    # x is a tensor of shape [n, m] while y is of shape [n].
    def objective(x, y):
      return tf.math.reduce_sum(x ** 2, axis=-1) + y ** 2
    ```
    The state in this case is specified by two input tensors `x` and `y`. To
    apply the algorithm to this objective function, one would need to specify
    either an initial population as a list of two tensors of shapes
    `[population_size, k]` and `[population_size]`. The following code shows the
    complete example:

    ```python
      population_size = 40
      # With an initial population and a multi-part state.
      initial_population = (tf.random.normal([population_size]),
                            tf.random.normal([population_size]))
      def easom_fn(x, y):
        return -(tf.math.cos(x) * tf.math.cos(y) *
                 tf.math.exp(-(x-np.pi)**2 - (y-np.pi)**2))

      optim_results = tfp.optimizers.differential_evolution_minimize(
          easom_fn,
          initial_population=initial_population,
          seed=43210)

      print (optim_results.converged)
      print (optim_results.position)  # Should be (close to) [pi, pi].
      print (optim_results.objective_value)    # Should be -1.

      # With a single starting point
      initial_position = (tf.constant(1.0), tf.constant(1.0))

      optim_results = tfp.optimizers.differential_evolution_minimize(
          easom_fn,
          initial_position=initial_position,
          population_size=40,
          population_stddev=2.0,
          seed=43210)
    ```

    Args:
      objective_function: A Python callable that accepts a batch of possible
        solutions and returns the values of the objective function at those
        arguments as a rank 1 real `Tensor`. This specifies the function to be
        minimized. The input to this callable may be either a single `Tensor`
        or a Python `list` of `Tensor`s. The signature must match the format of
        the argument `population`. (i.e. objective_function(*population) must
        return the value of the function to be minimized).
      initial_population: A real `Tensor` or Python list of `Tensor`s.
        If a list, each `Tensor` must be of rank at least 1 and with a common
        first dimension. The first dimension indexes into the candidate solutions
        while the rest of the dimensions (if any) index into an individual
        solution. The size of the population must be at least 4. This is a
        requirement of the DE algorithm.
      initial_position: A real `Tensor` of any shape. The seed solution used
        to initialize the population of solutions. If this parameter is specified
        then `initial_population` must not be specified.
      population_size: A positive scalar int32 `Tensor` greater than 4. The
        size of the population to evolve. This parameter is ignored if
        `initial_population` is specified.
        Default value: 50.
      population_stddev: A positive scalar real `Tensor` of the same dtype
        as `initial_position`. This parameter is ignored if `initial_population`
        is specified. Used to generate the population from the `initial_position`
        by adding random normal noise with zero mean and the specified standard
        deviation.
        Default value: 1.0
      max_iterations: Positive scalar int32 `Tensor`. The maximum number of
        generations to evolve the population for.
        Default value: 100
      func_tolerance: Scalar `Tensor` of the same dtype as the output of the
        `objective_function`. The algorithm stops if the absolute difference
        between the largest and the smallest objective function value in the
        population is below this number.
        Default value: 0
      position_tolerance: Scalar `Tensor` of the same real dtype as
        `initial_position` or `initial_population`. The algorithm terminates if
        the largest absolute difference between the coordinates of the population
        members is below this threshold.
        Default value: 1e-8
      differential_weight: Real scalar `Tensor`. Must be positive and less than
        2.0. The parameter controlling the strength of mutation in the algorithm.
        Default value: 0.5
      crossover_prob: Real scalar `Tensor`. Must be between 0 and 1. The
        probability of recombination per site.
        Default value: 0.9
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: None.
      name: (Optional) Python str. The name prefixed to the ops created by this
        function. If not supplied, the default name
        'differential_evolution_minimize' is used.
        Default value: None

    Returns:
      optimizer_results: An object containing the following attributes:
        converged: Scalar boolean `Tensor` indicating whether the minimum was
          found within the specified tolerances.
        num_objective_evaluations: The total number of objective
          evaluations performed.
        position: A `Tensor` containing the best point found during the search.
          If the search converged, then this value is the argmin of the
          objective function within the specified tolerances.
        objective_value: A `Tensor` containing the value of the objective
          function at the `position`. If the search
          converged, then this is the (local) minimum of
          the objective function.
        final_population: The final state of the population.
        final_objective_values: The objective function evaluated at the
          final population.
        initial_population: The starting population.
        initial_objective_values: The objective function evaluated at the
          initial population.
        num_iterations: The number of iterations of the main algorithm body.

    Raises:
      ValueError: If neither the initial population, nor the initial position
        are specified or if both are specified.


was_iterable = False<EOL>if initial_position is not None:<EOL><INDENT>initial_position, was_iterable = _ensure_list(initial_position)<EOL><DEDENT>if initial_population is not None:<EOL><INDENT>initial_population, was_iterable = _ensure_list(initial_population)<EOL><DEDENT>population = _get_starting_population(initial_population,<EOL>initial_position,<EOL>population_size,<EOL>population_stddev,<EOL>seed=seed)<EOL>differential_weight = tf.convert_to_tensor(<EOL>value=differential_weight, dtype=population[<NUM_LIT:0>].dtype.base_dtype)<EOL>crossover_prob = tf.convert_to_tensor(value=crossover_prob)<EOL>population_values = objective_function(*population)<EOL>if max_iterations is not None:<EOL><INDENT>max_iterations = tf.convert_to_tensor(value=max_iterations)<EOL><DEDENT>func_tolerance = tf.convert_to_tensor(<EOL>value=func_tolerance, dtype=population_values.dtype.base_dtype)<EOL>position_tolerance = tf.convert_to_tensor(<EOL>value=position_tolerance, dtype=population[<NUM_LIT:0>].dtype.base_dtype)<EOL>return (was_iterable,<EOL>population,<EOL>population_values,<EOL>max_iterations,<EOL>func_tolerance,<EOL>position_tolerance,<EOL>differential_weight,<EOL>crossover_prob)<EOL>


return tf.math.reduce_all(input_tensor=tf.math.is_inf(population_values))<EOL>

Checks if all the population values are NaN/infinite.


best_value = tf.math.reduce_min(input_tensor=values)<EOL>best_index = tf.where(tf.math.equal(values, best_value))[<NUM_LIT:0>, <NUM_LIT:0>]<EOL>return ([population_part[best_index] for population_part in population],<EOL>best_value)<EOL>

Finds the population member with the lowest value.


<EOL>value_range = tf.math.abs(<EOL>tf.math.reduce_max(input_tensor=population_values) -<EOL>tf.math.reduce_min(input_tensor=population_values))<EOL>value_converged = value_range <= func_tolerance<EOL>half_tol = position_tolerance / <NUM_LIT:2><EOL>def part_converged(part):<EOL><INDENT>return tf.math.reduce_max(input_tensor=tf.math.abs(part -<EOL>part[<NUM_LIT:0>])) <= half_tol<EOL><DEDENT>x_converged = tf.math.reduce_all(<EOL>input_tensor=[part_converged(part) for part in population])<EOL>return value_converged | x_converged<EOL>

Checks whether the convergence criteria have been met.


if initial_population is not None:<EOL><INDENT>return [tf.convert_to_tensor(value=part) for part in initial_population]<EOL><DEDENT>seed_stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>population = []<EOL>for part in initial_position:<EOL><INDENT>part = tf.convert_to_tensor(value=part)<EOL>part_event_shape = tf.shape(input=part)<EOL>population_part_shape = tf.concat([[population_size-<NUM_LIT:1>],<EOL>part_event_shape], axis=<NUM_LIT:0>)<EOL>population_part = tf.random.normal(population_part_shape,<EOL>stddev=population_stddev,<EOL>dtype=part.dtype.base_dtype,<EOL>seed=seed_stream())<EOL>population_part += part<EOL>population_part = tf.concat([[part], population_part], axis=<NUM_LIT:0>)<EOL>population.append(population_part)<EOL><DEDENT>return population<EOL>

Constructs the initial population.

    If an initial population is not already provided, this function constructs
    a population by adding random normal noise to the initial position.

    Args:
      initial_population: None or a list of `Tensor`s. The initial population.
      initial_position: None or a list of `Tensor`s. The initial position.
        If initial_population is None, this argument must not be None.
      population_size: Scalar integer `Tensor`. The number of members in the
        population. If the initial population is not None, this parameter is
        ignored.
      population_stddev: A positive scalar real `Tensor` of the same dtype
        as `initial_position` or `initial_population` (whichever is not None).
        This parameter is ignored if `initial_population`
        is specified. Used to generate the population from the
        `initial_position` by adding random normal noise with zero mean and
        the specified standard deviation.
      seed: Seed for random number generation.

    Returns:
      A list of `Tensor`s. The initial population.


sizes = [tf.cast(tf.size(input=x), dtype=tf.float64) for x in population]<EOL>seed_stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>force_crossover_group = distributions.Categorical(sizes).sample(<EOL>[population_size, <NUM_LIT:1>], seed=seed_stream())<EOL>recombinants = []<EOL>for i, population_part in enumerate(population):<EOL><INDENT>pop_part_flat = tf.reshape(population_part, [population_size, -<NUM_LIT:1>])<EOL>mutant_part_flat = tf.reshape(mutants[i], [population_size, -<NUM_LIT:1>])<EOL>part_size = tf.size(input=population_part) // population_size<EOL>force_crossovers = tf.one_hot(<EOL>tf.random.uniform([population_size],<EOL>minval=<NUM_LIT:0>,<EOL>maxval=part_size,<EOL>dtype=tf.int32,<EOL>seed=seed_stream()),<EOL>part_size,<EOL>on_value=True,<EOL>off_value=False,<EOL>dtype=tf.bool)  <EOL>group_mask = tf.math.equal(force_crossover_group, i)<EOL>force_crossovers &= group_mask<EOL>do_binary_crossover = tf.random.uniform(<EOL>[population_size, part_size],<EOL>dtype=crossover_prob.dtype.base_dtype,<EOL>seed=seed_stream()) < crossover_prob<EOL>do_binary_crossover |= force_crossovers<EOL>recombinant_flat = tf.where(<EOL>do_binary_crossover,<EOL>x=mutant_part_flat,<EOL>y=pop_part_flat)<EOL>recombinant = tf.reshape(recombinant_flat, tf.shape(input=population_part))<EOL>recombinants.append(recombinant)<EOL><DEDENT>return recombinants<EOL>

Performs recombination by binary crossover for the current population.

    Let v_i denote the i'th component of the member v and m_i the corresponding
    component of the mutant vector corresponding to v. Then the crossed over
    vector w_i is determined by setting w_i =
    (m_i with probability=crossover_prob else v_i). In addition, DE requires that
    at least one of the components is crossed over (otherwise we end
    up with no change). This is done by choosing on index say k randomly where
    a force crossover is performed (i.e. w_k = m_k). This is the scheme
    implemented in this function.

    Args:
      population: A Python list of `Tensor`s where each `Tensor` in the list
        must be of rank at least 1 and all the elements must have a common
        first dimension. The base population to cross over.
      population_size: A scalar integer `Tensor`. The number of elements in the
        population (i.e. size of the first dimension of any member of
        `population`).
      mutants: A Python list of `Tensor`s with the same structure as `population`.
        The mutated population.
      crossover_prob: A postive real scalar `Tensor` bounded above by 1.0. The
        probability of a crossover being performed for each axis.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.

    Returns:
      A list of `Tensor`s of the same structure, dtype and shape as `population`.
      The recombined population.


mixing_indices = tf.reshape(mixing_indices, [-<NUM_LIT:1>])<EOL>weights = tf.stack([<NUM_LIT:1.0>, differential_weight, -differential_weight])<EOL>def _mutant_part(population_part):<EOL><INDENT>donors = tf.gather(population_part, mixing_indices)<EOL>donors = tf.transpose(<EOL>a=tf.reshape(donors, [population_size, <NUM_LIT:3>, -<NUM_LIT:1>]), perm=[<NUM_LIT:0>, <NUM_LIT:2>, <NUM_LIT:1>])<EOL>return tf.math.reduce_sum(input_tensor=donors * weights, axis=-<NUM_LIT:1>)<EOL><DEDENT>return [_mutant_part(population_part) for population_part in population]<EOL>

Computes the mutatated vectors for each population member.

    Args:
      population:  Python `list` of `Tensor`s representing the
        current population vectors. Each `Tensor` must be of the same real dtype.
        The first dimension of each `Tensor` indexes individual
        population members. For example, if the population is a list with a
        single `Tensor` of shape [n, m1, m2], then `n` is the population size and
        the shape of an individual solution is [m1, m2].
        If there is more than one element in the population, then each `Tensor`
        in the list should have the first axis of the same size.
      population_size: Scalar integer `Tensor`. The size of the population.
      mixing_indices: `Tensor` of integral dtype and shape [n, 3] where `n` is the
        number of members in the population. Each element of the `Tensor` must be
        a valid index into the first dimension of the population (i.e range
        between `0` and `n-1` inclusive).
      differential_weight: Real scalar `Tensor`. Must be positive and less than
        2.0. The parameter controlling the strength of mutation.

    Returns:
      mutants: `Tensor` or Python `list` of `Tensor`s of the same shape and dtype
        as the input population. The mutated vectors.


def _get_mixing_indices(size, seed=None, name=None):

with tf.compat.v1.name_scope(<EOL>name, default_name='<STR_LIT>', values=[size]):<EOL><INDENT>size = tf.convert_to_tensor(value=size)<EOL>dtype = size.dtype<EOL>seed_stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>first = tf.random.uniform([size],<EOL>maxval=size-<NUM_LIT:1>,<EOL>dtype=dtype,<EOL>seed=seed_stream())<EOL>second = tf.random.uniform([size],<EOL>maxval=size-<NUM_LIT:2>,<EOL>dtype=dtype,<EOL>seed=seed_stream())<EOL>third = tf.random.uniform([size],<EOL>maxval=size-<NUM_LIT:3>,<EOL>dtype=dtype,<EOL>seed=seed_stream())<EOL>second = tf.where(first < second, x=second, y=second + <NUM_LIT:1>)<EOL>smaller = tf.math.minimum(first, second)<EOL>larger = tf.math.maximum(first, second)<EOL>third = tf.where(third < smaller, x=third, y=third + <NUM_LIT:1>)<EOL>third = tf.where(third < larger, x=third, y=third + <NUM_LIT:1>)<EOL>sample = tf.stack([first, second, third], axis=<NUM_LIT:1>)<EOL>to_avoid = tf.expand_dims(tf.range(size), axis=-<NUM_LIT:1>)<EOL>sample = tf.where(sample < to_avoid, x=sample, y=sample + <NUM_LIT:1>)<EOL>return sample<EOL><DEDENT>

Generates an array of indices suitable for mutation operation.

    The mutation operation in differential evolution requires that for every
    element of the population, three distinct other elements be chosen to produce
    a trial candidate. This function generates an array of shape [size, 3]
    satisfying the properties that:
      (a). array[i, :] does not contain the index 'i'.
      (b). array[i, :] does not contain any overlapping indices.
      (c). All elements in the array are between 0 and size - 1 inclusive.

    Args:
      size: Scalar integer `Tensor`. The number of samples as well as a the range
        of the indices to sample from.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.
      name:  Python `str` name prefixed to Ops created by this function.
        Default value: 'get_mixing_indices'.

    Returns:
      sample: A `Tensor` of shape [size, 3] and same dtype as `size` containing
      samples without replacement between 0 and size - 1 (inclusive) with the
      `i`th row not including the number `i`.


if isinstance(tensor_or_list, (list, tuple)):<EOL><INDENT>return list(tensor_or_list), True<EOL><DEDENT>return [tensor_or_list], False<EOL>

Converts the input arg to a list if it is not a list already.

    Args:
      tensor_or_list: A `Tensor` or a Python list of `Tensor`s. The argument to
        convert to a list of `Tensor`s.

    Returns:
      A tuple of two elements. The first is a Python list of `Tensor`s containing
      the original arguments. The second is a boolean indicating whether
      the original argument was a list or tuple already.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[initial_position, tolerance, initial_inverse_hessian_estimate]):<EOL><INDENT>initial_position = tf.convert_to_tensor(<EOL>value=initial_position, name='<STR_LIT>')<EOL>dtype = initial_position.dtype.base_dtype<EOL>tolerance = tf.convert_to_tensor(<EOL>value=tolerance, dtype=dtype, name='<STR_LIT>')<EOL>f_relative_tolerance = tf.convert_to_tensor(<EOL>value=f_relative_tolerance, dtype=dtype, name='<STR_LIT>')<EOL>x_tolerance = tf.convert_to_tensor(<EOL>value=x_tolerance, dtype=dtype, name='<STR_LIT>')<EOL>max_iterations = tf.convert_to_tensor(<EOL>value=max_iterations, name='<STR_LIT>')<EOL>input_shape = distribution_util.prefer_static_shape(initial_position)<EOL>batch_shape, domain_size = input_shape[:-<NUM_LIT:1>], input_shape[-<NUM_LIT:1>]<EOL>if stopping_condition is None:<EOL><INDENT>stopping_condition = bfgs_utils.converged_all<EOL><DEDENT>control_inputs = None<EOL>if initial_inverse_hessian_estimate is None:<EOL><INDENT>initial_inv_hessian = tf.eye(domain_size,<EOL>batch_shape=batch_shape,<EOL>dtype=dtype,<EOL>name='<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>initial_inv_hessian = tf.convert_to_tensor(<EOL>value=initial_inverse_hessian_estimate,<EOL>dtype=dtype,<EOL>name='<STR_LIT>')<EOL>control_inputs = _inv_hessian_control_inputs(initial_inv_hessian)<EOL>hessian_shape = tf.concat([batch_shape, [domain_size, domain_size]], <NUM_LIT:0>)<EOL>initial_inv_hessian = tf.broadcast_to(initial_inv_hessian, hessian_shape)<EOL><DEDENT>def _cond(state):<EOL><INDENT>"""<STR_LIT>"""<EOL>return ((state.num_iterations < max_iterations) &<EOL>tf.logical_not(stopping_condition(state.converged, state.failed)))<EOL><DEDENT>def _body(state):<EOL><INDENT>"""<STR_LIT>"""<EOL>search_direction = _get_search_direction(state.inverse_hessian_estimate,<EOL>state.objective_gradient)<EOL>derivative_at_start_pt = tf.reduce_sum(<EOL>input_tensor=state.objective_gradient * search_direction, axis=-<NUM_LIT:1>)<EOL>needs_reset = (~state.failed & ~state.converged &<EOL>(derivative_at_start_pt >= <NUM_LIT:0>))<EOL>search_direction_reset = _get_search_direction(<EOL>initial_inv_hessian, state.objective_gradient)<EOL>actual_serch_direction = tf.where(<EOL>needs_reset, search_direction_reset, search_direction)<EOL>actual_inv_hessian = tf.where(<EOL>needs_reset, initial_inv_hessian, state.inverse_hessian_estimate)<EOL>current_state = bfgs_utils.update_fields(<EOL>state, inverse_hessian_estimate=actual_inv_hessian)<EOL>next_state = bfgs_utils.line_search_step(<EOL>current_state,<EOL>value_and_gradients_function, actual_serch_direction,<EOL>tolerance, f_relative_tolerance, x_tolerance, stopping_condition)<EOL>return [_update_inv_hessian(current_state, next_state)]<EOL><DEDENT>kwargs = bfgs_utils.get_initial_state_args(<EOL>value_and_gradients_function,<EOL>initial_position,<EOL>tolerance,<EOL>control_inputs)<EOL>kwargs['<STR_LIT>'] = initial_inv_hessian<EOL>initial_state = BfgsOptimizerResults(**kwargs)<EOL>return tf.while_loop(<EOL>cond=_cond,<EOL>body=_body,<EOL>loop_vars=[initial_state],<EOL>parallel_iterations=parallel_iterations)[<NUM_LIT:0>]<EOL><DEDENT>

Applies the BFGS algorithm to minimize a differentiable function.

    Performs unconstrained minimization of a differentiable function using the
    BFGS scheme. For details of the algorithm, see [Nocedal and Wright(2006)][1].

    ### Usage:

    The following example demonstrates the BFGS optimizer attempting to find the
    minimum for a simple two dimensional quadratic objective function.

    ```python
      minimum = np.array([1.0, 1.0])  # The center of the quadratic bowl.
      scales = np.array([2.0, 3.0])  # The scales along the two axes.

      # The objective function and the gradient.
      def quadratic(x):
        value = tf.reduce_sum(scales * (x - minimum) ** 2)
        return value, tf.gradients(value, x)[0]

      start = tf.constant([0.6, 0.8])  # Starting point for the search.
      optim_results = tfp.optimizer.bfgs_minimize(
          quadratic, initial_position=start, tolerance=1e-8)

      with tf.Session() as session:
        results = session.run(optim_results)
        # Check that the search converged
        assert(results.converged)
        # Check that the argmin is close to the actual value.
        np.testing.assert_allclose(results.position, minimum)
        # Print out the total number of function evaluations it took. Should be 6.
        print ("Function evaluations: %d" % results.num_objective_evaluations)
    ```

    ### References:
    [1]: Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series in
      Operations Research. pp 136-140. 2006
      http://pages.mtu.edu/~struther/Courses/OLD/Sp2013/5630/Jorge_Nocedal_Numerical_optimization_267490.pdf

    Args:
      value_and_gradients_function:  A Python callable that accepts a point as a
        real `Tensor` and returns a tuple of `Tensor`s of real dtype containing
        the value of the function and its gradient at that point. The function
        to be minimized. The input should be of shape `[..., n]`, where `n` is
        the size of the domain of input points, and all others are batching
        dimensions. The first component of the return value should be a real
        `Tensor` of matching shape `[...]`. The second component (the gradient)
        should also be of shape `[..., n]` like the input value to the function.
      initial_position: real `Tensor` of shape `[..., n]`. The starting point, or
        points when using batching dimensions, of the search procedure. At these
        points the function value and the gradient norm should be finite.
      tolerance: Scalar `Tensor` of real dtype. Specifies the gradient tolerance
        for the procedure. If the supremum norm of the gradient vector is below
        this number, the algorithm is stopped.
      x_tolerance: Scalar `Tensor` of real dtype. If the absolute change in the
        position between one iteration and the next is smaller than this number,
        the algorithm is stopped.
      f_relative_tolerance: Scalar `Tensor` of real dtype. If the relative change
        in the objective value between one iteration and the next is smaller
        than this value, the algorithm is stopped.
      initial_inverse_hessian_estimate: Optional `Tensor` of the same dtype
        as the components of the output of the `value_and_gradients_function`.
        If specified, the shape should broadcastable to shape `[..., n, n]`; e.g.
        if a single `[n, n]` matrix is provided, it will be automatically
        broadcasted to all batches. Alternatively, one can also specify a
        different hessian estimate for each batch member.
        For the correctness of the algorithm, it is required that this parameter
        be symmetric and positive definite. Specifies the starting estimate for
        the inverse of the Hessian at the initial point. If not specified,
        the identity matrix is used as the starting estimate for the
        inverse Hessian.
      max_iterations: Scalar positive int32 `Tensor`. The maximum number of
        iterations for BFGS updates.
      parallel_iterations: Positive integer. The number of iterations allowed to
        run in parallel.
      stopping_condition: (Optional) 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. The default is tfp.optimizer.converged_all
        which only stops when all batch members have either converged or failed.
        An alternative is tfp.optimizer.converged_any which stops as soon as one
        batch member has converged, or when all have failed.
      name: (Optional) Python str. The name prefixed to the ops created by this
        function. If not supplied, the default name 'minimize' is used.

    Returns:
      optimizer_results: A namedtuple containing the following items:
        converged: boolean tensor of shape `[...]` indicating for each batch
          member whether the minimum was found within tolerance.
        failed:  boolean tensor of shape `[...]` indicating for each batch
          member whether a line search step failed to find a suitable step size
          satisfying Wolfe conditions. In the absence of any constraints on the
          number of objective evaluations permitted, this value will
          be the complement of `converged`. However, if there is
          a constraint and the search stopped due to available
          evaluations being exhausted, both `failed` and `converged`
          will be simultaneously False.
        num_objective_evaluations: The total number of objective
          evaluations performed.
        position: A tensor of shape `[..., n]` containing the last argument value
          found during the search from each starting point. If the search
          converged, then this value is the argmin of the objective function.
        objective_value: A tensor of shape `[...]` with the value of the
          objective function at the `position`. If the search converged, then
          this is the (local) minimum of the objective function.
        objective_gradient: A tensor of shape `[..., n]` containing the gradient
          of the objective function at the `position`. If the search converged
          the max-norm of this tensor should be below the tolerance.
        inverse_hessian_estimate: A tensor of shape `[..., n, n]` containing the
          inverse of the estimated Hessian.


<EOL>is_positive_definite = tf.reduce_all(<EOL>input_tensor=tf.math.is_finite(tf.linalg.cholesky(inv_hessian)),<EOL>axis=[-<NUM_LIT:1>, -<NUM_LIT:2>])<EOL>is_symmetric = tf.equal(bfgs_utils.norm(<EOL>inv_hessian - _batch_transpose(inv_hessian), dims=<NUM_LIT:2>), <NUM_LIT:0>)<EOL>return [tf.Assert(is_positive_definite,<EOL>['<STR_LIT>',<EOL>inv_hessian]),<EOL>tf.Assert(is_symmetric,<EOL>['<STR_LIT>',<EOL>inv_hessian])]<EOL>

Computes control inputs to validate a provided inverse Hessian.

    These ensure that the provided inverse Hessian is positive definite and
    symmetric.

    Args:
      inv_hessian: The starting estimate for the inverse of the Hessian at the
        initial point.

    Returns:
      A list of tf.Assert ops suitable for use with tf.control_dependencies.


return -_mul_right(inv_hessian_approx, gradient)<EOL>

Computes the direction along which to perform line search.


<EOL>should_update = ~next_state.converged & ~next_state.failed<EOL>gradient_delta = next_state.objective_gradient - prev_state.objective_gradient<EOL>position_delta = next_state.position - prev_state.position<EOL>normalization_factor = tf.reduce_sum(<EOL>input_tensor=gradient_delta * position_delta, axis=-<NUM_LIT:1>)<EOL>should_update = should_update & ~tf.equal(normalization_factor, <NUM_LIT:0>)<EOL>def _do_update_inv_hessian():<EOL><INDENT>next_inv_hessian = _bfgs_inv_hessian_update(<EOL>gradient_delta, position_delta, normalization_factor,<EOL>prev_state.inverse_hessian_estimate)<EOL>return bfgs_utils.update_fields(<EOL>next_state,<EOL>inverse_hessian_estimate=tf.where(should_update,<EOL>next_inv_hessian,<EOL>prev_state.inverse_hessian_estimate))<EOL><DEDENT>return prefer_static.cond(<EOL>tf.reduce_any(input_tensor=should_update),<EOL>_do_update_inv_hessian,<EOL>lambda: next_state)<EOL>

Update the BGFS state by computing the next inverse hessian estimate.