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

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

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

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

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

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

def line_search_step ( state , value_and_gradients_function , search_direction , grad_tolerance , f_relative_tolerance , x_tolerance , stopping_condition ) : line_search_value_grad_func = _restrict_along_direction ( value_and_gradients_function , state . position , search_direction ) derivative_at_start_pt = tf . reduce_sum ( input_tensor = state . objective_gradient * search_direction , axis = - 1 ) val_0 = ValueAndGradient ( x = _broadcast ( 0 , state . position ) , f = state . objective_value , df = derivative_at_start_pt , full_gradient = state . objective_gradient ) inactive = state . failed | state . converged ls_result = linesearch . hager_zhang ( line_search_value_grad_func , initial_step_size = _broadcast ( 1 , state . position ) , value_at_zero = val_0 , converged = inactive ) state_after_ls = update_fields ( state , failed = state . failed | ~ ls_result . converged , num_iterations = state . num_iterations + 1 , num_objective_evaluations = ( state . num_objective_evaluations + ls_result . func_evals ) ) def _do_update_position ( ) : position_delta = tf . where ( inactive , tf . zeros_like ( search_direction ) , search_direction * tf . expand_dims ( ls_result . left . x , axis = - 1 ) ) return _update_position ( state_after_ls , position_delta , ls_result . left . f , ls_result . left . full_gradient , grad_tolerance , f_relative_tolerance , x_tolerance ) return prefer_static . cond ( stopping_condition ( state . converged , state . failed ) , true_fn = lambda : state_after_ls , false_fn = _do_update_position )

Performs the line search step of the BFGS search procedure .


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

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

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

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

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

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

  return _restricted_func

def _restrict_along_direction ( value_and_gradients_function , position , direction ) : def _restricted_func ( t ) : t = _broadcast ( t , position ) pt = position + tf . expand_dims ( t , axis = - 1 ) * direction objective_value , gradient = value_and_gradients_function ( pt ) return ValueAndGradient ( x = t , f = objective_value , df = tf . reduce_sum ( input_tensor = gradient * direction , axis = - 1 ) , full_gradient = gradient ) return _restricted_func

Restricts a function in n - dimensions to a given direction .


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

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

def _update_position ( state , position_delta , next_objective , next_gradient , grad_tolerance , f_relative_tolerance , x_tolerance ) : failed = state . failed | ~ tf . math . is_finite ( next_objective ) | ~ tf . reduce_all ( input_tensor = tf . math . is_finite ( next_gradient ) , axis = - 1 ) next_position = state . position + position_delta converged = ~ failed & _check_convergence ( state . position , next_position , state . objective_value , next_objective , next_gradient , grad_tolerance , f_relative_tolerance , x_tolerance ) return update_fields ( state , converged = state . converged | converged , failed = failed , position = next_position , objective_value = next_objective , objective_gradient = next_gradient )

Updates the state advancing its position by a given position_delta .


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

Checks if the algorithm satisfies the convergence criteria .


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

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

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

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

def _broadcast ( value , target ) : return tf . broadcast_to ( tf . convert_to_tensor ( value = value , dtype = target . dtype ) , distribution_util . prefer_static_shape ( target ) [ : - 1 ] )

Broadcast a value to match the batching dimensions of a target .


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

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

  Args:
    x: input float.

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

def _harmonic_number ( x ) : one = tf . ones ( [ ] , dtype = x . dtype ) return tf . math . digamma ( x + one ) - tf . math . digamma ( one )

Compute the harmonic number from its analytic continuation .


def default_exchange_proposed_fn(prob_exchange):
  """Default exchange proposal function, for replica exchange MC.

  With probability `prob_exchange` propose combinations of replica for exchange.
  When exchanging, create combinations of adjacent replicas in
  [Replica Exchange Monte Carlo](
  https://en.wikipedia.org/wiki/Parallel_tempering)

  ```
  exchange_fn = default_exchange_proposed_fn(prob_exchange=0.5)
  exchange_proposed = exchange_fn(num_replica=3)

  exchange_proposed.eval()
  ==> [[0, 1]]  # 1 exchange, 0 <--> 1

  exchange_proposed.eval()
  ==> []  # 0 exchanges
  ```

  Args:
    prob_exchange: Scalar `Tensor` giving probability that any exchanges will
      be generated.

  Returns:
    default_exchange_proposed_fn_: Python callable which take a number of
      replicas (a Python integer), and return combinations of replicas for
      exchange as an [n, 2] integer `Tensor`, `0 <= n <= num_replica // 2`,
      with *unique* values in the set `{0, ..., num_replica}`.
  """

  def default_exchange_proposed_fn_(num_replica, seed=None):
    """Default function for `exchange_proposed_fn` of `kernel`."""
    seed_stream = distributions.SeedStream(seed, 'default_exchange_proposed_fn')

    zero_start = tf.random.uniform([], seed=seed_stream()) > 0.5
    if num_replica % 2 == 0:

      def _exchange():
        flat_exchange = tf.range(num_replica)
        if num_replica > 2:
          start = tf.cast(~zero_start, dtype=tf.int32)
          end = num_replica - start
          flat_exchange = flat_exchange[start:end]
        return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // 2, 2])
    else:

      def _exchange():
        start = tf.cast(zero_start, dtype=tf.int32)
        end = num_replica - tf.cast(~zero_start, dtype=tf.int32)
        flat_exchange = tf.range(num_replica)[start:end]
        return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // 2, 2])

    def _null_exchange():
      return tf.reshape(tf.cast([], dtype=tf.int32), shape=[0, 2])

    return tf.cond(
        pred=tf.random.uniform([], seed=seed_stream()) < prob_exchange,
        true_fn=_exchange,
        false_fn=_null_exchange)

  return default_exchange_proposed_fn_

def default_exchange_proposed_fn ( prob_exchange ) : def default_exchange_proposed_fn_ ( num_replica , seed = None ) : seed_stream = distributions . SeedStream ( seed , 'default_exchange_proposed_fn' ) zero_start = tf . random . uniform ( [ ] , seed = seed_stream ( ) ) > 0.5 if num_replica % 2 == 0 : def _exchange ( ) : flat_exchange = tf . range ( num_replica ) if num_replica > 2 : start = tf . cast ( ~ zero_start , dtype = tf . int32 ) end = num_replica - start flat_exchange = flat_exchange [ start : end ] return tf . reshape ( flat_exchange , [ tf . size ( input = flat_exchange ) // 2 , 2 ] ) else : def _exchange ( ) : start = tf . cast ( zero_start , dtype = tf . int32 ) end = num_replica - tf . cast ( ~ zero_start , dtype = tf . int32 ) flat_exchange = tf . range ( num_replica ) [ start : end ] return tf . reshape ( flat_exchange , [ tf . size ( input = flat_exchange ) // 2 , 2 ] ) def _null_exchange ( ) : return tf . reshape ( tf . cast ( [ ] , dtype = tf . int32 ) , shape = [ 0 , 2 ] ) return tf . cond ( pred = tf . random . uniform ( [ ] , seed = seed_stream ( ) ) < prob_exchange , true_fn = _exchange , false_fn = _null_exchange ) return default_exchange_proposed_fn_

Default exchange proposal function for replica exchange MC .


def _get_field(kernel_results, field_name):
  """field_name from kernel_results or kernel_results.accepted_results."""
  if hasattr(kernel_results, field_name):
    return getattr(kernel_results, field_name)
  if hasattr(kernel_results, 'accepted_results'):
    return getattr(kernel_results.accepted_results, field_name)
  raise TypeError('Cannot extract %s from %s' % (field_name, kernel_results))

def
_get_field
(
kernel_results
,
field_name
)
:
if
hasattr
(
kernel_results
,
field_name
)
:
return
getattr
(
kernel_results
,
field_name
)
if
hasattr
(
kernel_results
,
'accepted_results'
)
:
return
getattr
(
kernel_results
.
accepted_results
,
field_name
)
raise
TypeError
(
'Cannot extract %s from %s'
%
(
field_name
,
kernel_results
)
)


def _get_exchanged_states(self, old_states, exchange_proposed,
                            exchange_proposed_n, sampled_replica_states,
                            sampled_replica_results):
    """Get list of TensorArrays holding exchanged states, and zeros."""
    with tf.compat.v1.name_scope('get_exchanged_states'):

      target_log_probs = []
      for replica in range(self.num_replica):
        replica_log_prob = _get_field(sampled_replica_results[replica],
                                      'target_log_prob')
        inverse_temp = self.inverse_temperatures[replica]
        target_log_probs.append(replica_log_prob / inverse_temp)
      target_log_probs = tf.stack(target_log_probs, axis=0)

      dtype = target_log_probs.dtype
      num_state_parts = len(sampled_replica_states[0])
      # exchanged_states[k][i] is Tensor of (new) state part k, for replica i.
      # The `k` will be known statically, and `i` is a Tensor.
      # We will insert values into indices `i` for every replica with a proposed
      # exchange.
      exchanged_states = [
          tf.TensorArray(
              dtype,
              size=self.num_replica,
              dynamic_size=False,
              tensor_array_name='exchanged_states',
              # State part k has same shape, regardless of replica.  So use 0.
              element_shape=sampled_replica_states[0][k].shape)
          for k in range(num_state_parts)
      ]

      # Draw random variables here, to avoid sampling in the loop (and losing
      # reproducibility).  This may mean we sample too many, but we will always
      # have enough.
      sample_shape = tf.concat(
          ([self.num_replica // 2], tf.shape(input=target_log_probs)[1:]),
          axis=0)
      log_uniforms = tf.math.log(
          tf.random.uniform(
              shape=sample_shape, dtype=dtype, seed=self._seed_stream()))

      def _swap(is_exchange_accepted, x, y):
        """Swap batches of x, y where accepted."""
        with tf.compat.v1.name_scope('swap_where_exchange_accepted'):
          new_x = mcmc_util.choose(is_exchange_accepted, y, x)
          new_y = mcmc_util.choose(is_exchange_accepted, x, y)
        return new_x, new_y

      def cond(i, unused_exchanged_states):
        return i < exchange_proposed_n

      def body(i, exchanged_states):
        """Body of while loop for exchanging states."""
        # Propose exchange between replicas indexed by m and n.
        m, n = tf.unstack(exchange_proposed[i])

        # Construct log_accept_ratio:  -temp_diff * target_log_prob_diff.
        # Note target_log_prob_diff = -EnergyDiff (common definition is in terms
        # of energy).
        temp_diff = self.inverse_temperatures[m] - self.inverse_temperatures[n]
        # Difference of target log probs may be +- Inf or NaN.  We want the
        # product of this with the temperature difference to have "alt value" of
        # -Inf.
        log_accept_ratio = mcmc_util.safe_sum(
            [-temp_diff * target_log_probs[m], temp_diff * target_log_probs[n]])

        is_exchange_accepted = log_uniforms[i] < log_accept_ratio

        for k in range(num_state_parts):
          new_m, new_n = _swap(is_exchange_accepted, old_states[k].read(m),
                               old_states[k].read(n))
          exchanged_states[k] = exchanged_states[k].write(m, new_m)
          exchanged_states[k] = exchanged_states[k].write(n, new_n)

        return i + 1, exchanged_states

      # At this point, exchanged_states[k] is a length num_replicas TensorArray.
      return tf.while_loop(
          cond=cond, body=body, loop_vars=[tf.constant(0),
                                           exchanged_states])[1]

def _get_exchanged_states ( self , old_states , exchange_proposed , exchange_proposed_n , sampled_replica_states , sampled_replica_results ) : with tf . compat . v1 . name_scope ( 'get_exchanged_states' ) : target_log_probs = [ ] for replica in range ( self . num_replica ) : replica_log_prob = _get_field ( sampled_replica_results [ replica ] , 'target_log_prob' ) inverse_temp = self . inverse_temperatures [ replica ] target_log_probs . append ( replica_log_prob / inverse_temp ) target_log_probs = tf . stack ( target_log_probs , axis = 0 ) dtype = target_log_probs . dtype num_state_parts = len ( sampled_replica_states [ 0 ] ) exchanged_states = [ tf . TensorArray ( dtype , size = self . num_replica , dynamic_size = False , tensor_array_name = 'exchanged_states' , element_shape = sampled_replica_states [ 0 ] [ k ] . shape ) for k in range ( num_state_parts ) ] sample_shape = tf . concat ( ( [ self . num_replica // 2 ] , tf . shape ( input = target_log_probs ) [ 1 : ] ) , axis = 0 ) log_uniforms = tf . math . log ( tf . random . uniform ( shape = sample_shape , dtype = dtype , seed = self . _seed_stream ( ) ) ) def _swap ( is_exchange_accepted , x , y ) : with tf . compat . v1 . name_scope ( 'swap_where_exchange_accepted' ) : new_x = mcmc_util . choose ( is_exchange_accepted , y , x ) new_y = mcmc_util . choose ( is_exchange_accepted , x , y ) return new_x , new_y def cond ( i , unused_exchanged_states ) : return i < exchange_proposed_n def body ( i , exchanged_states ) : m , n = tf . unstack ( exchange_proposed [ i ] ) temp_diff = self . inverse_temperatures [ m ] - self . inverse_temperatures [ n ] log_accept_ratio = mcmc_util . safe_sum ( [ - temp_diff * target_log_probs [ m ] , temp_diff * target_log_probs [ n ] ] ) is_exchange_accepted = log_uniforms [ i ] < log_accept_ratio for k in range ( num_state_parts ) : new_m , new_n = _swap ( is_exchange_accepted , old_states [ k ] . read ( m ) , old_states [ k ] . read ( n ) ) exchanged_states [ k ] = exchanged_states [ k ] . write ( m , new_m ) exchanged_states [ k ] = exchanged_states [ k ] . write ( n , new_n ) return i + 1 , exchanged_states return tf . while_loop ( cond = cond , body = body , loop_vars = [ tf . constant ( 0 ) , exchanged_states ] ) [ 1 ]

Get list of TensorArrays holding exchanged states and zeros .


def _variance_scale_term(self):
    """Helper to `_covariance` and `_variance` which computes a shared scale."""
    # Expand back the last dim so the shape of _variance_scale_term matches the
    # shape of self.concentration.
    c0 = self.total_concentration[..., tf.newaxis]
    return tf.sqrt((1. + c0 / self.total_count[..., tf.newaxis]) / (1. + c0))

def _variance_scale_term ( self ) : c0 = self . total_concentration [ ... , tf . newaxis ] return tf . sqrt ( ( 1. + c0 / self . total_count [ ... , tf . newaxis ] ) / ( 1. + c0 ) )

Helper to _covariance and _variance which computes a shared scale .


def forward_log_det_jacobian_fn(bijector):
  """Makes a function which applies a list of Bijectors' `log_det_jacobian`s."""
  if not mcmc_util.is_list_like(bijector):
    bijector = [bijector]

  def fn(transformed_state_parts, event_ndims):
    return sum([
        b.forward_log_det_jacobian(sp, event_ndims=e)
        for b, e, sp in zip(bijector, event_ndims, transformed_state_parts)
    ])

  return fn

def forward_log_det_jacobian_fn ( bijector ) : if not mcmc_util . is_list_like ( bijector ) : bijector = [ bijector ] def fn ( transformed_state_parts , event_ndims ) : return sum ( [ b . forward_log_det_jacobian ( sp , event_ndims = e ) for b , e , sp in zip ( bijector , event_ndims , transformed_state_parts ) ] ) return fn

Makes a function which applies a list of Bijectors log_det_jacobian s .


def forward_transform_fn(bijector):
  """Makes a function which applies a list of Bijectors' `forward`s."""
  if not mcmc_util.is_list_like(bijector):
    bijector = [bijector]

  def fn(transformed_state_parts):
    return [b.forward(sp) for b, sp in zip(bijector, transformed_state_parts)]

  return fn

def forward_transform_fn ( bijector ) : if not mcmc_util . is_list_like ( bijector ) : bijector = [ bijector ] def fn ( transformed_state_parts ) : return [ b . forward ( sp ) for b , sp in zip ( bijector , transformed_state_parts ) ] return fn

Makes a function which applies a list of Bijectors forward s .


def inverse_transform_fn(bijector):
  """Makes a function which applies a list of Bijectors' `inverse`s."""
  if not mcmc_util.is_list_like(bijector):
    bijector = [bijector]
  def fn(state_parts):
    return [b.inverse(sp)
            for b, sp in zip(bijector, state_parts)]
  return fn

def inverse_transform_fn ( bijector ) : if not mcmc_util . is_list_like ( bijector ) : bijector = [ bijector ] def fn ( state_parts ) : return [ b . inverse ( sp ) for b , sp in zip ( bijector , state_parts ) ] return fn

Makes a function which applies a list of Bijectors inverse s .


def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of the Transformed Kernel.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s
        representing the current state(s) of the Markov chain(s),
        _after_ application of `bijector.forward`. The first `r`
        dimensions index independent chains,
        `r = tf.rank(target_log_prob_fn(*current_state))`. The
        `inner_kernel.one_step` does not actually use `current_state`,
        rather it takes as input
        `previous_kernel_results.transformed_state` (because
        `TransformedTransitionKernel` creates a copy of the input
        inner_kernel with a modified `target_log_prob_fn` which
        internally applies the `bijector.forward`).
      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.
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'transformed_kernel', 'one_step'),
        values=[previous_kernel_results]):
      transformed_next_state, kernel_results = self._inner_kernel.one_step(
          previous_kernel_results.transformed_state,
          previous_kernel_results.inner_results)
      transformed_next_state_parts = (
          transformed_next_state
          if mcmc_util.is_list_like(transformed_next_state) else
          [transformed_next_state])
      next_state_parts = self._forward_transform(transformed_next_state_parts)
      next_state = (
          next_state_parts if mcmc_util.is_list_like(transformed_next_state)
          else next_state_parts[0])
      kernel_results = TransformedTransitionKernelResults(
          transformed_state=transformed_next_state,
          inner_results=kernel_results)
      return next_state, kernel_results

Runs one iteration of the Transformed Kernel .


def val_where(cond, tval, fval):
  """Like tf.where but works on namedtuples."""
  if isinstance(tval, tf.Tensor):
    return tf.where(cond, tval, fval)
  elif isinstance(tval, tuple):
    cls = type(tval)
    return cls(*(val_where(cond, t, f) for t, f in zip(tval, fval)))
  else:
    raise Exception(TypeError)

def val_where ( cond , tval , fval ) : if isinstance ( tval , tf . Tensor ) : return tf . where ( cond , tval , fval ) elif isinstance ( tval , tuple ) : cls = type ( tval ) return cls ( * ( val_where ( cond , t , f ) for t , f in zip ( tval , fval ) ) ) else : raise Exception ( TypeError )

Like tf . where but works on namedtuples .


def secant2(value_and_gradients_function,
            val_0,
            search_interval,
            f_lim,
            sufficient_decrease_param=0.1,
            curvature_param=0.9,
            name=None):
  """Performs the secant square procedure of Hager Zhang.

  Given an interval that brackets a root, this procedure performs an update of
  both end points using two intermediate points generated using the secant
  interpolation. For details see the steps S1-S4 in [Hager and Zhang (2006)][2].

  The interval [a, b] must satisfy the opposite slope conditions described in
  the documentation for `update`.

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns an object that can be converted to a namedtuple.
      The namedtuple should have fields 'f' and 'df' that correspond to scalar
      tensors of real dtype containing the value of the function and its
      derivative at that point. The other namedtuple fields, if present,
      should be tensors or sequences (possibly nested) of tensors.
      In usual optimization application, this function would be generated by
      projecting the multivariate objective function along some specific
      direction. The direction is determined by some other procedure but should
      be a descent direction (i.e. the derivative of the projected univariate
      function must be negative at 0.).
      Alternatively, the function may represent the batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and the fields 'f' and 'df' in the returned
      namedtuple should each be a tensor of shape [n], with the corresponding
      function values and derivatives at the input points.
    val_0: A namedtuple, as returned by value_and_gradients_function evaluated
      at `0.`.
    search_interval: A namedtuple describing the current search interval,
      must include the fields:
      - converged: Boolean `Tensor` of shape [n], indicating batch members
          where search has already converged. Interval for these batch members
          won't be modified.
      - failed: Boolean `Tensor` of shape [n], indicating batch members
          where search has already failed. Interval for these batch members
          wont be modified.
      - iterations: Scalar int32 `Tensor`. Number of line search iterations
          so far.
      - func_evals: Scalar int32 `Tensor`. Number of function evaluations
          so far.
      - left: A namedtuple, as returned by value_and_gradients_function,
          of the left end point of the current search interval.
      - right: A namedtuple, as returned by value_and_gradients_function,
          of the right end point of the current search interval.
    f_lim: Scalar `Tensor` of real dtype. The function value threshold for
      the approximate Wolfe conditions to be checked.
    sufficient_decrease_param: Positive scalar `Tensor` of real dtype.
      Bounded above by the curvature param. Corresponds to 'delta' in the
      terminology of [Hager and Zhang (2006)][2].
    curvature_param: Positive scalar `Tensor` of real dtype. Bounded above
      by `1.`. Corresponds to 'sigma' in the terminology of
      [Hager and Zhang (2006)][2].
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'secant2' is used.

  Returns:
    A namedtuple containing the following fields.
      active: A boolean `Tensor` of shape [n]. Used internally by the procedure
        to indicate batch members on which there is work left to do.
      converged: A boolean `Tensor` of shape [n]. Indicates whether a point
        satisfying the Wolfe conditions has been found. If this is True, the
        interval will be degenerate (i.e. `left` and `right` below will be
        identical).
      failed: A boolean `Tensor` of shape [n]. Indicates if invalid function or
        gradient values were encountered (i.e. infinity or NaNs).
      num_evals: A scalar int32 `Tensor`. The total number of function
        evaluations made.
      left: Return value of value_and_gradients_function at the updated left
        end point of the interval.
      right: Return value of value_and_gradients_function at the updated right
        end point of the interval.
  """
  with tf.compat.v1.name_scope(name, 'secant2', [
      val_0, search_interval, f_lim, sufficient_decrease_param,
      curvature_param]):
    # This will always be s.t. left <= c <= right
    val_c = value_and_gradients_function(
        _secant(search_interval.left, search_interval.right))
    failed = search_interval.failed | ~is_finite(val_c)
    converged = search_interval.converged | (~failed & _satisfies_wolfe(
        val_0, val_c, f_lim, sufficient_decrease_param, curvature_param))
    new_converged = converged & ~search_interval.converged
    val_left = val_where(new_converged, val_c, search_interval.left)
    val_right = val_where(new_converged, val_c, search_interval.right)

    initial_args = _Secant2Result(
        active=~failed & ~converged,
        converged=converged,
        failed=failed,
        num_evals=search_interval.func_evals + 1,
        left=val_left,
        right=val_right)

    def _apply_secant2_inner():
      return _secant2_inner(
          value_and_gradients_function,
          initial_args,
          val_0,
          val_c,
          f_lim,
          sufficient_decrease_param,
          curvature_param)

    return prefer_static.cond(
        tf.reduce_any(input_tensor=initial_args.active),
        _apply_secant2_inner,
        lambda: initial_args)

def secant2 ( value_and_gradients_function , val_0 , search_interval , f_lim , sufficient_decrease_param = 0.1 , curvature_param = 0.9 , name = None ) : with tf . compat . v1 . name_scope ( name , 'secant2' , [ val_0 , search_interval , f_lim , sufficient_decrease_param , curvature_param ] ) : val_c = value_and_gradients_function ( _secant ( search_interval . left , search_interval . right ) ) failed = search_interval . failed | ~ is_finite ( val_c ) converged = search_interval . converged | ( ~ failed & _satisfies_wolfe ( val_0 , val_c , f_lim , sufficient_decrease_param , curvature_param ) ) new_converged = converged & ~ search_interval . converged val_left = val_where ( new_converged , val_c , search_interval . left ) val_right = val_where ( new_converged , val_c , search_interval . right ) initial_args = _Secant2Result ( active = ~ failed & ~ converged , converged = converged , failed = failed , num_evals = search_interval . func_evals + 1 , left = val_left , right = val_right ) def _apply_secant2_inner ( ) : return _secant2_inner ( value_and_gradients_function , initial_args , val_0 , val_c , f_lim , sufficient_decrease_param , curvature_param ) return prefer_static . cond ( tf . reduce_any ( input_tensor = initial_args . active ) , _apply_secant2_inner , lambda : initial_args )

Performs the secant square procedure of Hager Zhang .


def _secant2_inner(value_and_gradients_function,
                   initial_args,
                   val_0,
                   val_c,
                   f_lim,
                   sufficient_decrease_param,
                   curvature_param):
  """Helper function for secant square."""
  # Apply the `update` function on active branch members to squeeze their
  # bracketing interval.
  update_result = update(value_and_gradients_function,
                         initial_args.left,
                         initial_args.right,
                         val_c,
                         f_lim,
                         active=initial_args.active)

  # Update active and failed flags, update left/right on non-failed entries.
  active = initial_args.active & ~update_result.failed
  failed = initial_args.failed | update_result.failed
  val_left = val_where(active, update_result.left, initial_args.left)
  val_right = val_where(active, update_result.right, initial_args.right)

  # Check if new `c` points should be generated.
  updated_left = active & tf.equal(val_left.x, val_c.x)
  updated_right = active & tf.equal(val_right.x, val_c.x)
  is_new = updated_left | updated_right

  next_c = tf.where(updated_left,
                    _secant(initial_args.left, val_left),
                    val_c.x)
  next_c = tf.where(updated_right,
                    _secant(initial_args.right, val_right),
                    next_c)
  in_range = (val_left.x <= next_c) & (next_c <= val_right.x)

  # Figure out if an extra function evaluation is needed for new `c` points.
  needs_extra_eval = tf.reduce_any(input_tensor=in_range & is_new)
  num_evals = initial_args.num_evals + update_result.num_evals
  num_evals = num_evals + tf.cast(needs_extra_eval, num_evals.dtype)

  next_args = _Secant2Result(
      active=active & in_range,  # No longer active if `c` is out of range.
      converged=initial_args.converged,
      failed=failed,
      num_evals=num_evals,
      left=val_left,
      right=val_right)

  def _apply_inner_update():
    next_val_c = prefer_static.cond(
        needs_extra_eval,
        (lambda: value_and_gradients_function(next_c)),
        (lambda: val_c))
    return _secant2_inner_update(
        value_and_gradients_function, next_args, val_0, next_val_c, f_lim,
        sufficient_decrease_param, curvature_param)

  return prefer_static.cond(
      tf.reduce_any(input_tensor=next_args.active),
      _apply_inner_update,
      lambda: next_args)

def _secant2_inner ( value_and_gradients_function , initial_args , val_0 , val_c , f_lim , sufficient_decrease_param , curvature_param ) : update_result = update ( value_and_gradients_function , initial_args . left , initial_args . right , val_c , f_lim , active = initial_args . active ) active = initial_args . active & ~ update_result . failed failed = initial_args . failed | update_result . failed val_left = val_where ( active , update_result . left , initial_args . left ) val_right = val_where ( active , update_result . right , initial_args . right ) updated_left = active & tf . equal ( val_left . x , val_c . x ) updated_right = active & tf . equal ( val_right . x , val_c . x ) is_new = updated_left | updated_right next_c = tf . where ( updated_left , _secant ( initial_args . left , val_left ) , val_c . x ) next_c = tf . where ( updated_right , _secant ( initial_args . right , val_right ) , next_c ) in_range = ( val_left . x <= next_c ) & ( next_c <= val_right . x ) needs_extra_eval = tf . reduce_any ( input_tensor = in_range & is_new ) num_evals = initial_args . num_evals + update_result . num_evals num_evals = num_evals + tf . cast ( needs_extra_eval , num_evals . dtype ) next_args = _Secant2Result ( active = active & in_range , converged = initial_args . converged , failed = failed , num_evals = num_evals , left = val_left , right = val_right ) def _apply_inner_update ( ) : next_val_c = prefer_static . cond ( needs_extra_eval , ( lambda : value_and_gradients_function ( next_c ) ) , ( lambda : val_c ) ) return _secant2_inner_update ( value_and_gradients_function , next_args , val_0 , next_val_c , f_lim , sufficient_decrease_param , curvature_param ) return prefer_static . cond ( tf . reduce_any ( input_tensor = next_args . active ) , _apply_inner_update , lambda : next_args )


def _secant2_inner_update(value_and_gradients_function,
                          initial_args,
                          val_0,
                          val_c,
                          f_lim,
                          sufficient_decrease_param,
                          curvature_param):
  """Helper function for secant-square step."""
  # Fail if `val_c` is no longer finite.
  new_failed = initial_args.active & ~is_finite(val_c)
  active = initial_args.active & ~new_failed
  failed = initial_args.failed | new_failed

  # We converge when we find a point satisfying the Wolfe conditions, in those
  # cases we set `val_left = val_right = val_c`.
  found_wolfe = active & _satisfies_wolfe(
      val_0, val_c, f_lim, sufficient_decrease_param, curvature_param)
  val_left = val_where(found_wolfe, val_c, initial_args.left)
  val_right = val_where(found_wolfe, val_c, initial_args.right)
  converged = initial_args.converged | found_wolfe
  active = active & ~found_wolfe

  # If any active batch members remain, we apply the `update` function to
  # squeeze further their corresponding left/right bracketing interval.
  def _apply_update():
    update_result = update(
        value_and_gradients_function, val_left, val_right, val_c, f_lim,
        active=active)
    return _Secant2Result(
        active=tf.zeros_like(active),  # End of secant2, no actives anymore.
        converged=converged,
        failed=failed | update_result.failed,
        num_evals=initial_args.num_evals + update_result.num_evals,
        left=update_result.left,
        right=update_result.right)

  # Otherwise just return the current results.
  def _default():
    return _Secant2Result(
        active=active,
        converged=converged,
        failed=failed,
        num_evals=initial_args.num_evals,
        left=val_left,
        right=val_right)

  return prefer_static.cond(
      tf.reduce_any(input_tensor=active), _apply_update, _default)

def _secant2_inner_update ( value_and_gradients_function , initial_args , val_0 , val_c , f_lim , sufficient_decrease_param , curvature_param ) : new_failed = initial_args . active & ~ is_finite ( val_c ) active = initial_args . active & ~ new_failed failed = initial_args . failed | new_failed found_wolfe = active & _satisfies_wolfe ( val_0 , val_c , f_lim , sufficient_decrease_param , curvature_param ) val_left = val_where ( found_wolfe , val_c , initial_args . left ) val_right = val_where ( found_wolfe , val_c , initial_args . right ) converged = initial_args . converged | found_wolfe active = active & ~ found_wolfe def _apply_update ( ) : update_result = update ( value_and_gradients_function , val_left , val_right , val_c , f_lim , active = active ) return _Secant2Result ( active = tf . zeros_like ( active ) , converged = converged , failed = failed | update_result . failed , num_evals = initial_args . num_evals + update_result . num_evals , left = update_result . left , right = update_result . right ) def _default ( ) : return _Secant2Result ( active = active , converged = converged , failed = failed , num_evals = initial_args . num_evals , left = val_left , right = val_right ) return prefer_static . cond ( tf . reduce_any ( input_tensor = active ) , _apply_update , _default )

Helper function for secant - square step .


def update(value_and_gradients_function, val_left, val_right, val_trial, f_lim,
           active=None):
  """Squeezes a bracketing interval containing the minimum.

  Given an interval which brackets a minimum and a point in that interval,
  finds a smaller nested interval which also brackets the minimum. If the
  supplied point does not lie in the bracketing interval, the current interval
  is returned.

  The following description is given in terms of individual points evaluated on
  a line function to be minimized. Note, however, the implementation also
  accepts batches of points allowing to minimize multiple line functions at
  once. See details on the docstring of `value_and_gradients_function` below.

  The requirement of the interval bracketing a minimum is expressed through the
  opposite slope conditions. Assume the left end point is 'a', the right
  end point is 'b', the function to be minimized is 'f' and the derivative is
  'df'. The update procedure relies on the following conditions being satisfied:

  '''
    f(a) <= f(0) + epsilon   (1)
    df(a) < 0                (2)
    df(b) > 0                (3)
  '''

  In the first condition, epsilon is a small positive constant. The condition
  demands that the function at the left end point be not much bigger than the
  starting point (i.e. 0). This is an easy to satisfy condition because by
  assumption, we are in a direction where the function value is decreasing.
  The second and third conditions together demand that there is at least one
  zero of the derivative in between a and b.

  In addition to the interval, the update algorithm requires a third point to
  be supplied. Usually, this point would lie within the interval [a, b]. If the
  point is outside this interval, the current interval is returned. If the
  point lies within the interval, the behaviour of the function and derivative
  value at this point is used to squeeze the original interval in a manner that
  preserves the opposite slope conditions.

  For further details of this component, see the procedure U0-U3 on page 123 of
  the [Hager and Zhang (2006)][2] article.

  Note that this function does not explicitly verify whether the opposite slope
  conditions are satisfied for the supplied interval. It is assumed that this
  is so.

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns an object that can be converted to a namedtuple.
      The namedtuple should have fields 'f' and 'df' that correspond to scalar
      tensors of real dtype containing the value of the function and its
      derivative at that point. The other namedtuple fields, if present,
      should be tensors or sequences (possibly nested) of tensors.
      In usual optimization application, this function would be generated by
      projecting the multivariate objective function along some specific
      direction. The direction is determined by some other procedure but should
      be a descent direction (i.e. the derivative of the projected univariate
      function must be negative at 0.).
      Alternatively, the function may represent the batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and the fields 'f' and 'df' in the returned
      namedtuple should each be a tensor of shape [n], with the corresponding
      function values and derivatives at the input points.
    val_left: Return value of value_and_gradients_function at the left
      end point of the bracketing interval (labelles 'a' above).
    val_right: Return value of value_and_gradients_function at the right
      end point of the bracketing interval (labelles 'b' above).
    val_trial: Return value of value_and_gradients_function at the trial point
      to be used to shrink the interval (labelled 'c' above).
    f_lim: real `Tensor` of shape [n]. The function value threshold for
      the approximate Wolfe conditions to be checked for each batch member.
    active: optional boolean `Tensor` of shape [n]. Relevant in batching mode
      only, indicates batch members on which the update procedure should be
      applied. On non-active members the current left/right interval is returned
      unmodified.

  Returns:
    A namedtuple containing the following fields:
      iteration: An int32 scalar `Tensor`. The number of iterations performed
        by the bisect algorithm.
      stopped: A boolean `Tensor` of shape [n]. True for those batch members
        where the bisection algorithm terminated.
      failed: A boolean `Tensor` of shape [n]. True for those batch members
        where an error was encountered.
      num_evals: An int32 scalar `Tensor`. The number of times the objective
        function was evaluated.
      left: Return value of value_and_gradients_function at the updated left
        end point of the interval found.
      right: Return value of value_and_gradients_function at the updated right
        end point of the interval found.
  """
  # We should only update if the trial point is within the interval.
  within_range = (val_left.x < val_trial.x) & (val_trial.x < val_right.x)
  if active is not None:
    within_range = within_range & active

  # The new point is a valid left end point if it has negative slope
  # and the value at the point is not too large.
  valid_left = (val_trial.df < 0) & (val_trial.f <= f_lim)

  # If the trial point has a negative slope but the value at that point
  # is too high, bisect can narrow down an interval between the current left
  # and the trial point.
  needs_bisect = within_range & (val_trial.df < 0) & (val_trial.f > f_lim)

  # Note that if `~valid_left` it is because either:
  # - the slope at the trial point is positive, so it is a valid right
  #   point, or
  # - the needs_bisect condition is true.
  # In both cases we want to keep the current left and replace right
  # with the trial point.
  left = val_where(within_range & valid_left, val_trial, val_left)
  right = val_where(within_range & ~valid_left, val_trial, val_right)

  bisect_args = _IntermediateResult(
      iteration=tf.convert_to_tensor(value=0),
      stopped=~needs_bisect,
      failed=tf.zeros_like(within_range),  # i.e. all false.
      num_evals=tf.convert_to_tensor(value=0),
      left=left,
      right=right)
  return _bisect(value_and_gradients_function, bisect_args, f_lim)

def update ( value_and_gradients_function , val_left , val_right , val_trial , f_lim , active = None ) : within_range = ( val_left . x < val_trial . x ) & ( val_trial . x < val_right . x ) if active is not None : within_range = within_range & active valid_left = ( val_trial . df < 0 ) & ( val_trial . f <= f_lim ) needs_bisect = within_range & ( val_trial . df < 0 ) & ( val_trial . f > f_lim ) left = val_where ( within_range & valid_left , val_trial , val_left ) right = val_where ( within_range & ~ valid_left , val_trial , val_right ) bisect_args = _IntermediateResult ( iteration = tf . convert_to_tensor ( value = 0 ) , stopped = ~ needs_bisect , failed = tf . zeros_like ( within_range ) , num_evals = tf . convert_to_tensor ( value = 0 ) , left = left , right = right ) return _bisect ( value_and_gradients_function , bisect_args , f_lim )

Squeezes a bracketing interval containing the minimum .


def bracket(value_and_gradients_function,
            search_interval,
            f_lim,
            max_iterations,
            expansion_param=5.0):
  """Brackets the minimum given an initial starting point.

  Applies the Hager Zhang bracketing algorithm to find an interval containing
  a region with points satisfying Wolfe conditions. Uses the supplied initial
  step size 'c', the right end point of the provided search interval, to find
  such an interval. The only condition on 'c' is that it should be positive.
  For more details see steps B0-B3 in [Hager and Zhang (2006)][2].

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns a namedtuple containing the value filed `f` of the
      function and its derivative value field `df` at that point.
      Alternatively, the function may representthe batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and return a tuple of two tensors of shape [n], the
      function values and the corresponding derivatives at the input points.
    search_interval: A namedtuple describing the current search interval,
      must include the fields:
      - converged: Boolean `Tensor` of shape [n], indicating batch members
          where search has already converged. Interval for these batch members
          wont be modified.
      - failed: Boolean `Tensor` of shape [n], indicating batch members
          where search has already failed. Interval for these batch members
          wont be modified.
      - iterations: Scalar int32 `Tensor`. Number of line search iterations
          so far.
      - func_evals: Scalar int32 `Tensor`. Number of function evaluations
          so far.
      - left: A namedtuple, as returned by value_and_gradients_function
          evaluated at 0, the left end point of the current interval.
      - right: A namedtuple, as returned by value_and_gradients_function,
          of the right end point of the current interval (labelled 'c' above).
    f_lim: real `Tensor` of shape [n]. The function value threshold for
      the approximate Wolfe conditions to be checked for each batch member.
    max_iterations: Int32 scalar `Tensor`. The maximum number of iterations
      permitted. The limit applies equally to all batch members.
    expansion_param: Scalar positive `Tensor` of real dtype. Must be greater
      than `1.`. Used to expand the initial interval in case it does not bracket
      a minimum.

  Returns:
    A namedtuple with the following fields.
      iteration: An int32 scalar `Tensor`. The number of iterations performed.
        Bounded above by `max_iterations` parameter.
      stopped: A boolean `Tensor` of shape [n]. True for those batch members
        where the algorithm terminated before reaching `max_iterations`.
      failed: A boolean `Tensor` of shape [n]. True for those batch members
        where an error was encountered during bracketing.
      num_evals: An int32 scalar `Tensor`. The number of times the objective
        function was evaluated.
      left: Return value of value_and_gradients_function at the updated left
        end point of the interval found.
      right: Return value of value_and_gradients_function at the updated right
        end point of the interval found.
  """
  already_stopped = search_interval.failed | search_interval.converged

  # If the slope at right end point is positive, step B1 in [2], then the given
  # initial points already bracket a minimum.
  bracketed = search_interval.right.df >= 0

  # Bisection is needed, step B2, if right end point almost works as a new left
  # end point but the objective value is too high.
  needs_bisect = (
      search_interval.right.df < 0) & (search_interval.right.f > f_lim)

  # In these three cases bracketing is already `stopped` and there is no need
  # to perform further evaluations. Otherwise the bracketing loop is needed to
  # expand the interval, step B3, until the conditions are met.
  initial_args = _IntermediateResult(
      iteration=search_interval.iterations,
      stopped=already_stopped | bracketed | needs_bisect,
      failed=search_interval.failed,
      num_evals=search_interval.func_evals,
      left=search_interval.left,
      right=search_interval.right)

  def _loop_cond(curr):
    return (curr.iteration <
            max_iterations) & ~tf.reduce_all(input_tensor=curr.stopped)

  def _loop_body(curr):
    """Main body of bracketing loop."""
    # The loop maintains the invariant that curr.stopped is true if we have
    # either: failed, successfully bracketed, or not yet bracketed but needs
    # bisect. On the only remaining case, step B3 in [2]. case we need to
    # expand and update the left/right values appropriately.
    new_right = value_and_gradients_function(expansion_param * curr.right.x)
    left = val_where(curr.stopped, curr.left, curr.right)
    right = val_where(curr.stopped, curr.right, new_right)

    # Updated the failed, bracketed, and needs_bisect conditions.
    failed = curr.failed | ~is_finite(right)
    bracketed = right.df >= 0
    needs_bisect = (right.df < 0) & (right.f > f_lim)
    return [_IntermediateResult(
        iteration=curr.iteration + 1,
        stopped=curr.stopped | failed | bracketed | needs_bisect,
        failed=failed,
        num_evals=curr.num_evals + 1,
        left=left,
        right=right)]

  bracket_result = tf.while_loop(
      cond=_loop_cond, body=_loop_body, loop_vars=[initial_args])[0]

  # For entries where bisect is still needed, mark them as not yet stopped,
  # reset the left end point, and run `_bisect` on them.
  needs_bisect = (
      (bracket_result.right.df < 0) & (bracket_result.right.f > f_lim))
  stopped = already_stopped | bracket_result.failed | ~needs_bisect
  left = val_where(stopped, bracket_result.left, search_interval.left)
  bisect_args = bracket_result._replace(stopped=stopped, left=left)
  return _bisect(value_and_gradients_function, bisect_args, f_lim)

def bracket ( value_and_gradients_function , search_interval , f_lim , max_iterations , expansion_param = 5.0 ) : already_stopped = search_interval . failed | search_interval . converged bracketed = search_interval . right . df >= 0 needs_bisect = ( search_interval . right . df < 0 ) & ( search_interval . right . f > f_lim ) initial_args = _IntermediateResult ( iteration = search_interval . iterations , stopped = already_stopped | bracketed | needs_bisect , failed = search_interval . failed , num_evals = search_interval . func_evals , left = search_interval . left , right = search_interval . right ) def _loop_cond ( curr ) : return ( curr . iteration < max_iterations ) & ~ tf . reduce_all ( input_tensor = curr . stopped ) def _loop_body ( curr ) : new_right = value_and_gradients_function ( expansion_param * curr . right . x ) left = val_where ( curr . stopped , curr . left , curr . right ) right = val_where ( curr . stopped , curr . right , new_right ) failed = curr . failed | ~ is_finite ( right ) bracketed = right . df >= 0 needs_bisect = ( right . df < 0 ) & ( right . f > f_lim ) return [ _IntermediateResult ( iteration = curr . iteration + 1 , stopped = curr . stopped | failed | bracketed | needs_bisect , failed = failed , num_evals = curr . num_evals + 1 , left = left , right = right ) ] bracket_result = tf . while_loop ( cond = _loop_cond , body = _loop_body , loop_vars = [ initial_args ] ) [ 0 ] needs_bisect = ( ( bracket_result . right . df < 0 ) & ( bracket_result . right . f > f_lim ) ) stopped = already_stopped | bracket_result . failed | ~ needs_bisect left = val_where ( stopped , bracket_result . left , search_interval . left ) bisect_args = bracket_result . _replace ( stopped = stopped , left = left ) return _bisect ( value_and_gradients_function , bisect_args , f_lim )

Brackets the minimum given an initial starting point .


def bisect(value_and_gradients_function,
           initial_left,
           initial_right,
           f_lim):
  """Bisects an interval and updates to satisfy opposite slope conditions.

  Corresponds to the step U3 in [Hager and Zhang (2006)][2].

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns a namedtuple containing the value filed `f` of the
      function and its derivative value field `df` at that point.
      Alternatively, the function may representthe batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and return a tuple of two tensors of shape [n], the
      function values and the corresponding derivatives at the input points.
    initial_left: Return value of value_and_gradients_function at the left end
      point of the current bracketing interval.
    initial_right: Return value of value_and_gradients_function at the right end
      point of the current bracketing interval.
    f_lim: real `Tensor` of shape [n]. The function value threshold for
      the approximate Wolfe conditions to be checked for each batch member.

  Returns:
    A namedtuple containing the following fields:
      iteration: An int32 scalar `Tensor`. The number of iterations performed.
        Bounded above by `max_iterations` parameter.
      stopped: A boolean scalar `Tensor`. True if the bisect algorithm
        terminated.
      failed: A scalar boolean tensor. Indicates whether the objective function
        failed to produce a finite value.
      num_evals: A scalar int32 tensor. The number of value and gradients
        function evaluations.
      left: Return value of value_and_gradients_function at the left end
        point of the bracketing interval found.
      right: Return value of value_and_gradients_function at the right end
        point of the bracketing interval found.
  """
  failed = ~is_finite(initial_left, initial_right)
  needs_bisect = (initial_right.df < 0) & (initial_right.f > f_lim)
  bisect_args = _IntermediateResult(
      iteration=tf.convert_to_tensor(value=0),
      stopped=failed | ~needs_bisect,
      failed=failed,
      num_evals=tf.convert_to_tensor(value=0),
      left=initial_left,
      right=initial_right)
  return _bisect(value_and_gradients_function, bisect_args, f_lim)

def bisect ( value_and_gradients_function , initial_left , initial_right , f_lim ) : failed = ~ is_finite ( initial_left , initial_right ) needs_bisect = ( initial_right . df < 0 ) & ( initial_right . f > f_lim ) bisect_args = _IntermediateResult ( iteration = tf . convert_to_tensor ( value = 0 ) , stopped = failed | ~ needs_bisect , failed = failed , num_evals = tf . convert_to_tensor ( value = 0 ) , left = initial_left , right = initial_right ) return _bisect ( value_and_gradients_function , bisect_args , f_lim )

Bisects an interval and updates to satisfy opposite slope conditions .


def _bisect(value_and_gradients_function, initial_args, f_lim):
  """Actual implementation of bisect given initial_args in a _BracketResult."""
  def _loop_cond(curr):
    # TODO(b/112524024): Also take into account max_iterations.
    return ~tf.reduce_all(input_tensor=curr.stopped)

  def _loop_body(curr):
    """Narrow down interval to satisfy opposite slope conditions."""
    mid = value_and_gradients_function((curr.left.x + curr.right.x) / 2)

    # Fail if function values at mid point are no longer finite; or left/right
    # points are so close to it that we can't distinguish them any more.
    failed = (curr.failed | ~is_finite(mid) |
              tf.equal(mid.x, curr.left.x) | tf.equal(mid.x, curr.right.x))

    # If mid point has a negative slope and the function value at that point is
    # small enough, we can use it as a new left end point to narrow down the
    # interval. If mid point has a positive slope, then we have found a suitable
    # right end point to bracket a minima within opposite slopes. Otherwise, the
    # mid point has a negative slope but the function value at that point is too
    # high to work as left end point, we are in the same situation in which we
    # started the loop so we just update the right end point and continue.
    to_update = ~(curr.stopped | failed)
    update_left = (mid.df < 0) & (mid.f <= f_lim)
    left = val_where(to_update & update_left, mid, curr.left)
    right = val_where(to_update & ~update_left, mid, curr.right)

    # We're done when the right end point has a positive slope.
    stopped = curr.stopped | failed | (right.df >= 0)

    return [_IntermediateResult(
        iteration=curr.iteration,
        stopped=stopped,
        failed=failed,
        num_evals=curr.num_evals + 1,
        left=left,
        right=right)]

  # The interval needs updating if the right end point has a negative slope and
  # the value of the function at that point is too high. It is not a valid left
  # end point but along with the current left end point, it encloses another
  # minima. The loop above tries to narrow the interval so that it satisfies the
  # opposite slope conditions.
  return tf.while_loop(
      cond=_loop_cond, body=_loop_body, loop_vars=[initial_args])[0]

def _bisect ( value_and_gradients_function , initial_args , f_lim ) : def _loop_cond ( curr ) : return ~ tf . reduce_all ( input_tensor = curr . stopped ) def _loop_body ( curr ) : mid = value_and_gradients_function ( ( curr . left . x + curr . right . x ) / 2 ) failed = ( curr . failed | ~ is_finite ( mid ) | tf . equal ( mid . x , curr . left . x ) | tf . equal ( mid . x , curr . right . x ) ) to_update = ~ ( curr . stopped | failed ) update_left = ( mid . df < 0 ) & ( mid . f <= f_lim ) left = val_where ( to_update & update_left , mid , curr . left ) right = val_where ( to_update & ~ update_left , mid , curr . right ) stopped = curr . stopped | failed | ( right . df >= 0 ) return [ _IntermediateResult ( iteration = curr . iteration , stopped = stopped , failed = failed , num_evals = curr . num_evals + 1 , left = left , right = right ) ] return tf . while_loop ( cond = _loop_cond , body = _loop_body , loop_vars = [ initial_args ] ) [ 0 ]

Actual implementation of bisect given initial_args in a _BracketResult .


def is_finite(val_1, val_2=None):
  """Checks if the supplied values are finite.

  Args:
    val_1: A namedtuple instance with the function value and derivative,
      as returned e.g. by value_and_gradients_function evaluations.
    val_2: (Optional) A namedtuple instance with the function value and
      derivative, as returned e.g. by value_and_gradients_function evaluations.

  Returns:
    is_finite: Scalar boolean `Tensor` indicating whether the function value
      and the derivative in `val_1` (and optionally in `val_2`) are all finite.
  """
  val_1_finite = tf.math.is_finite(val_1.f) & tf.math.is_finite(val_1.df)
  if val_2 is not None:
    return val_1_finite & tf.math.is_finite(val_2.f) & tf.math.is_finite(
        val_2.df)
  return val_1_finite

def is_finite ( val_1 , val_2 = None ) : val_1_finite = tf . math . is_finite ( val_1 . f ) & tf . math . is_finite ( val_1 . df ) if val_2 is not None : return val_1_finite & tf . math . is_finite ( val_2 . f ) & tf . math . is_finite ( val_2 . df ) return val_1_finite

Checks if the supplied values are finite .


def _satisfies_wolfe(val_0,
                     val_c,
                     f_lim,
                     sufficient_decrease_param,
                     curvature_param):
  """Checks whether the Wolfe or approx Wolfe conditions are satisfied.

  The Wolfe conditions are a set of stopping criteria for an inexact line search
  algorithm. Let f(a) be the function value along the search direction and
  df(a) the derivative along the search direction evaluated a distance 'a'.
  Here 'a' is the distance along the search direction. The Wolfe conditions are:

    ```None
      f(a) <= f(0) + delta * a * df(0)   (Armijo/Sufficient decrease condition)
      df(a) >= sigma * df(0)             (Weak curvature condition)
    ```
  `delta` and `sigma` are two user supplied parameters satisfying:
   `0 < delta < sigma <= 1.`. In the following, delta is called
   `sufficient_decrease_param` and sigma is called `curvature_param`.

  On a finite precision machine, the Wolfe conditions are difficult to satisfy
  when one is close to the minimum. Hence, Hager-Zhang propose replacing
  the sufficient decrease condition with the following condition on the
  derivative in the vicinity of a minimum.

    ```None
      df(a) <= (2 * delta - 1) * df(0)  (Approx Wolfe sufficient decrease)
    ```
  This condition is only used if one is near the minimum. This is tested using

    ```None
      f(a) <= f(0) + epsilon * |f(0)|
    ```
  The following function checks both the Wolfe and approx Wolfe conditions.
  Here, `epsilon` is a small positive constant. In the following, the argument
  `f_lim` corresponds to the product: epsilon * |f(0)|.

  Args:
    val_0: A namedtuple, as returned by value_and_gradients_function
      evaluated at 0.
    val_c: A namedtuple, as returned by value_and_gradients_function
      evaluated at the point to be tested.
    f_lim: Scalar `Tensor` of real dtype. The function value threshold for
      the approximate Wolfe conditions to be checked.
    sufficient_decrease_param: Positive scalar `Tensor` of real dtype.
      Bounded above by the curvature param. Corresponds to 'delta' in the
      terminology of [Hager and Zhang (2006)][2].
    curvature_param: Positive scalar `Tensor` of real dtype. Bounded above
      by `1.`. Corresponds to 'sigma' in the terminology of
      [Hager Zhang (2005)][1].

  Returns:
    is_satisfied: A scalar boolean `Tensor` which is True if either the
      Wolfe or approximate Wolfe conditions are satisfied.
  """
  exact_wolfe_suff_dec = (sufficient_decrease_param * val_0.df >=
                          (val_c.f - val_0.f) / val_c.x)
  wolfe_curvature = val_c.df >= curvature_param * val_0.df
  exact_wolfe = exact_wolfe_suff_dec & wolfe_curvature
  approx_wolfe_applies = val_c.f <= f_lim
  approx_wolfe_suff_dec = ((2 * sufficient_decrease_param - 1) * val_0.df
                           >= val_c.df)
  approx_wolfe = approx_wolfe_applies & approx_wolfe_suff_dec & wolfe_curvature
  is_satisfied = exact_wolfe | approx_wolfe
  return is_satisfied

Checks whether the Wolfe or approx Wolfe conditions are satisfied .


def _secant(val_a, val_b):
  """Returns the secant interpolation for the minimum.

  The secant method is a technique for finding roots of nonlinear functions.
  When finding the minimum, one applies the secant method to the derivative
  of the function.
  For an arbitrary function and a bounding interval, the secant approximation
  can produce the next point which is outside the bounding interval. However,
  with the assumption of opposite slope condtion on the interval [a,b] the new
  point c is always bracketed by [a,b]. Note that by assumption,
  f'(a) < 0 and f'(b) > 0.
  Hence c is a weighted average of a and b and thus always in [a, b].

  Args:
    val_a: A namedtuple with the left end point, function value and derivative,
      of the current interval (i.e. a).
    val_b: A namedtuple with the right end point, function value and derivative,
      of the current interval (i.e. b).

  Returns:
    approx_minimum: A scalar real `Tensor`. An approximation to the point
      at which the derivative vanishes.
  """
  return (val_a.x * val_b.df - val_b.x * val_a.df) / (val_b.df - val_a.df)

def _secant ( val_a , val_b ) : return ( val_a . x * val_b . df - val_b . x * val_a . df ) / ( val_b . df - val_a . df )

Returns the secant interpolation for the minimum .


def make_simple_step_size_update_policy(num_adaptation_steps,
                                        target_rate=0.75,
                                        decrement_multiplier=0.01,
                                        increment_multiplier=0.01,
                                        step_counter=None):
  """Create a function implementing a step-size update policy.

  The simple policy increases or decreases the `step_size_var` based on the
  average of `exp(minimum(0., log_accept_ratio))`. It is based on
  [Section 4.2 of Andrieu and Thoms (2008)](
  https://people.eecs.berkeley.edu/~jordan/sail/readings/andrieu-thoms.pdf).

  The `num_adaptation_steps` argument is set independently of any burnin
  for the overall chain. In general, adaptation prevents the chain from
  reaching a stationary distribution, so obtaining consistent samples requires
  `num_adaptation_steps` be set to a value [somewhat smaller](
  http://andrewgelman.com/2017/12/15/burn-vs-warm-iterative-simulation-algorithms/#comment-627745)
  than the number of burnin steps. However, it may sometimes be helpful to set
  `num_adaptation_steps` to a larger value during development in order to
  inspect the behavior of the chain during adaptation.

  Args:
    num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
      during which to adjust the step size. This may be greater, less than, or
      equal to the number of burnin steps. If `None`, the step size is adapted
      on every step (note this breaks stationarity of the chain!).
    target_rate: Scalar `Tensor` representing desired `accept_ratio`.
      Default value: `0.75` (i.e., [center of asymptotically optimal
      rate](https://arxiv.org/abs/1411.6669)).
    decrement_multiplier: `Tensor` representing amount to downscale current
      `step_size`.
      Default value: `0.01`.
    increment_multiplier: `Tensor` representing amount to upscale current
      `step_size`.
      Default value: `0.01`.
    step_counter: Scalar `int` `Variable` specifying the current step. The step
      size is adapted iff `step_counter < num_adaptation_steps`.
      Default value: if `None`, an internal variable
        `step_size_adaptation_step_counter` is created and initialized to `-1`.

  Returns:
    step_size_simple_update_fn: Callable that takes args
      `step_size_var, kernel_results` and returns updated step size(s).
  """
  if step_counter is None and num_adaptation_steps is not None:
    step_counter = tf.compat.v1.get_variable(
        name='step_size_adaptation_step_counter',
        initializer=np.array(-1, dtype=np.int32),
        # Specify the dtype for variable sharing to work correctly
        # (b/120599991).
        dtype=tf.int32,
        trainable=False,
        use_resource=True)

  def step_size_simple_update_fn(step_size_var, kernel_results):
    """Updates (list of) `step_size` using a standard adaptive MCMC procedure.

    Args:
      step_size_var: (List of) `tf.Variable`s representing the per `state_part`
        HMC `step_size`.
      kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from most recent call to `one_step`.

    Returns:
      step_size_assign: (List of) `Tensor`(s) representing updated
        `step_size_var`(s).
    """

    if kernel_results is None:
      if mcmc_util.is_list_like(step_size_var):
        return [tf.identity(ss) for ss in step_size_var]
      return tf.identity(step_size_var)
    log_n = tf.math.log(
        tf.cast(
            tf.size(input=kernel_results.log_accept_ratio),
            kernel_results.log_accept_ratio.dtype))
    log_mean_accept_ratio = tf.reduce_logsumexp(
        input_tensor=tf.minimum(kernel_results.log_accept_ratio, 0.)) - log_n
    adjustment = tf.where(
        log_mean_accept_ratio < tf.cast(
            tf.math.log(target_rate), log_mean_accept_ratio.dtype),
        -decrement_multiplier / (1. + decrement_multiplier),
        increment_multiplier)

    def build_assign_op():
      if mcmc_util.is_list_like(step_size_var):
        return [
            ss.assign_add(ss * tf.cast(adjustment, ss.dtype))
            for ss in step_size_var
        ]
      return step_size_var.assign_add(
          step_size_var * tf.cast(adjustment, step_size_var.dtype))

    if num_adaptation_steps is None:
      return build_assign_op()
    else:
      with tf.control_dependencies([step_counter.assign_add(1)]):
        return tf.cond(
            pred=step_counter < num_adaptation_steps,
            true_fn=build_assign_op,
            false_fn=lambda: step_size_var)

  return step_size_simple_update_fn

def make_simple_step_size_update_policy ( num_adaptation_steps , target_rate = 0.75 , decrement_multiplier = 0.01 , increment_multiplier = 0.01 , step_counter = None ) : if step_counter is None and num_adaptation_steps is not None : step_counter = tf . compat . v1 . get_variable ( name = 'step_size_adaptation_step_counter' , initializer = np . array ( - 1 , dtype = np . int32 ) , dtype = tf . int32 , trainable = False , use_resource = True ) def step_size_simple_update_fn ( step_size_var , kernel_results ) : if kernel_results is None : if mcmc_util . is_list_like ( step_size_var ) : return [ tf . identity ( ss ) for ss in step_size_var ] return tf . identity ( step_size_var ) log_n = tf . math . log ( tf . cast ( tf . size ( input = kernel_results . log_accept_ratio ) , kernel_results . log_accept_ratio . dtype ) ) log_mean_accept_ratio = tf . reduce_logsumexp ( input_tensor = tf . minimum ( kernel_results . log_accept_ratio , 0. ) ) - log_n adjustment = tf . where ( log_mean_accept_ratio < tf . cast ( tf . math . log ( target_rate ) , log_mean_accept_ratio . dtype ) , - decrement_multiplier / ( 1. + decrement_multiplier ) , increment_multiplier ) def build_assign_op ( ) : if mcmc_util . is_list_like ( step_size_var ) : return [ ss . assign_add ( ss * tf . cast ( adjustment , ss . dtype ) ) for ss in step_size_var ] return step_size_var . assign_add ( step_size_var * tf . cast ( adjustment , step_size_var . dtype ) ) if num_adaptation_steps is None : return build_assign_op ( ) else : with tf . control_dependencies ( [ step_counter . assign_add ( 1 ) ] ) : return tf . cond ( pred = step_counter < num_adaptation_steps , true_fn = build_assign_op , false_fn = lambda : step_size_var ) return step_size_simple_update_fn

Create a function implementing a step - size update policy .


def _leapfrog_integrator_one_step(
    target_log_prob_fn,
    independent_chain_ndims,
    step_sizes,
    current_momentum_parts,
    current_state_parts,
    current_target_log_prob,
    current_target_log_prob_grad_parts,
    state_gradients_are_stopped=False,
    name=None):
  """Applies `num_leapfrog_steps` of the leapfrog integrator.

  Assumes a simple quadratic kinetic energy function: `0.5 ||momentum||**2`.

  #### Examples:

  ##### Simple quadratic potential.

  ```python
  import matplotlib.pyplot as plt
  %matplotlib inline
  import numpy as np
  import tensorflow as tf
  from tensorflow_probability.python.mcmc.hmc import _leapfrog_integrator_one_step  # pylint: disable=line-too-long
  tfd = tfp.distributions

  dims = 10
  num_iter = int(1e3)
  dtype = np.float32

  position = tf.placeholder(np.float32)
  momentum = tf.placeholder(np.float32)

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

  def _leapfrog_one_step(*args):
    # Closure representing computation done during each leapfrog step.
    return _leapfrog_integrator_one_step(
        target_log_prob_fn=target_log_prob_fn,
        independent_chain_ndims=0,
        step_sizes=[0.1],
        current_momentum_parts=args[0],
        current_state_parts=args[1],
        current_target_log_prob=args[2],
        current_target_log_prob_grad_parts=args[3])

  # Do leapfrog integration.
  [
      [next_momentum],
      [next_position],
      next_target_log_prob,
      next_target_log_prob_grad_parts,
  ] = tf.while_loop(
      cond=lambda *args: True,
      body=_leapfrog_one_step,
      loop_vars=[
        [momentum],
        [position],
        target_log_prob_fn(position),
        tf.gradients(target_log_prob_fn(position), position),
      ],
      maximum_iterations=3)

  momentum_ = np.random.randn(dims).astype(dtype)
  position_ = np.random.randn(dims).astype(dtype)
  positions = np.zeros([num_iter, dims], dtype)

  with tf.Session() as sess:
    for i in xrange(num_iter):
      position_, momentum_ = sess.run(
          [next_momentum, next_position],
          feed_dict={position: position_, momentum: momentum_})
      positions[i] = position_

  plt.plot(positions[:, 0]);  # Sinusoidal.
  ```

  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.
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    step_sizes: Python `list` of `Tensor`s representing the step size for the
      leapfrog integrator. Must broadcast with the shape of
      `current_state_parts`.  Larger step sizes lead to faster progress, but
      too-large step sizes make rejection exponentially more likely. When
      possible, it's often helpful to match per-variable step sizes to the
      standard deviations of the target distribution in each variable.
    current_momentum_parts: Tensor containing the value(s) of the momentum
      variable(s) to update.
    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.
    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.
    current_target_log_prob_grad_parts: Python list of `Tensor`s representing
      gradient of `target_log_prob_fn(*current_state_parts`) wrt
      `current_state_parts`. Must have same shape as `current_state_parts`. The
      only reason to specify this argument is to reduce TF graph size.
    state_gradients_are_stopped: Python `bool` indicating that the proposed new
      state be run through `tf.stop_gradient`. This is particularly useful when
      combining optimization over samples from the HMC chain.
      Default value: `False` (i.e., do not apply `stop_gradient`).
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'hmc_leapfrog_integrator').

  Returns:
    proposed_momentum_parts: Updated value of the momentum.
    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`.
    proposed_target_log_prob_grad_parts: Gradient of `proposed_target_log_prob`
      wrt `next_state`.

  Raises:
    ValueError: if `len(momentum_parts) != len(state_parts)`.
    ValueError: if `len(state_parts) != len(step_sizes)`.
    ValueError: if `len(state_parts) != len(grads_target_log_prob)`.
    TypeError: if `not target_log_prob.dtype.is_floating`.
  """
  # Note on per-variable step sizes:
  #
  # Using per-variable step sizes is equivalent to using the same step
  # size for all variables and adding a diagonal mass matrix in the
  # kinetic energy term of the Hamiltonian being integrated. This is
  # hinted at by Neal (2011) but not derived in detail there.
  #
  # Let x and v be position and momentum variables respectively.
  # Let g(x) be the gradient of `target_log_prob_fn(x)`.
  # Let S be a diagonal matrix of per-variable step sizes.
  # Let the Hamiltonian H(x, v) = -target_log_prob_fn(x) + 0.5 * ||v||**2.
  #
  # Using per-variable step sizes gives the updates
  # v'  = v  + 0.5 * matmul(S, g(x))
  # x'' = x  + matmul(S, v')
  # v'' = v' + 0.5 * matmul(S, g(x''))
  #
  # Let u = matmul(inv(S), v).
  # Multiplying v by inv(S) in the updates above gives the transformed dynamics
  # u'  = matmul(inv(S), v')  = matmul(inv(S), v) + 0.5 * g(x)
  #                           = u + 0.5 * g(x)
  # x'' = x + matmul(S, v') = x + matmul(S**2, u')
  # u'' = matmul(inv(S), v'') = matmul(inv(S), v') + 0.5 * g(x'')
  #                           = u' + 0.5 * g(x'')
  #
  # These are exactly the leapfrog updates for the Hamiltonian
  # H'(x, u) = -target_log_prob_fn(x) + 0.5 * u^T S**2 u
  #          = -target_log_prob_fn(x) + 0.5 * ||v||**2 = H(x, v).
  #
  # To summarize:
  #
  # * Using per-variable step sizes implicitly simulates the dynamics
  #   of the Hamiltonian H' (which are energy-conserving in H'). We
  #   keep track of v instead of u, but the underlying dynamics are
  #   the same if we transform back.
  # * The value of the Hamiltonian H'(x, u) is the same as the value
  #   of the original Hamiltonian H(x, v) after we transform back from
  #   u to v.
  # * Sampling v ~ N(0, I) is equivalent to sampling u ~ N(0, S**-2).
  #
  # So using per-variable step sizes in HMC will give results that are
  # exactly identical to explicitly using a diagonal mass matrix.

  with tf.compat.v1.name_scope(name, 'hmc_leapfrog_integrator_one_step', [
      independent_chain_ndims, step_sizes, current_momentum_parts,
      current_state_parts, current_target_log_prob,
      current_target_log_prob_grad_parts
  ]):

    # Step 1: Update momentum.
    proposed_momentum_parts = [
        v + 0.5 * tf.cast(eps, v.dtype) * g
        for v, eps, g
        in zip(current_momentum_parts,
               step_sizes,
               current_target_log_prob_grad_parts)]

    # Step 2: Update state.
    proposed_state_parts = [
        x + tf.cast(eps, v.dtype) * v
        for x, eps, v
        in zip(current_state_parts,
               step_sizes,
               proposed_momentum_parts)]

    if state_gradients_are_stopped:
      proposed_state_parts = [tf.stop_gradient(x) for x in proposed_state_parts]

    # Step 3a: Re-evaluate target-log-prob (and grad) at proposed state.
    [
        proposed_target_log_prob,
        proposed_target_log_prob_grad_parts,
    ] = mcmc_util.maybe_call_fn_and_grads(
        target_log_prob_fn,
        proposed_state_parts)

    if not proposed_target_log_prob.dtype.is_floating:
      raise TypeError('`target_log_prob_fn` must produce a `Tensor` '
                      'with `float` `dtype`.')

    if any(g is None for g in proposed_target_log_prob_grad_parts):
      raise ValueError(
          'Encountered `None` gradient. Does your target `target_log_prob_fn` '
          'access all `tf.Variable`s via `tf.get_variable`?\n'
          '  current_state_parts: {}\n'
          '  proposed_state_parts: {}\n'
          '  proposed_target_log_prob_grad_parts: {}'.format(
              current_state_parts,
              proposed_state_parts,
              proposed_target_log_prob_grad_parts))

    # Step 3b: Update momentum (again).
    proposed_momentum_parts = [
        v + 0.5 * tf.cast(eps, v.dtype) * g
        for v, eps, g
        in zip(proposed_momentum_parts,
               step_sizes,
               proposed_target_log_prob_grad_parts)]

    return [
        proposed_momentum_parts,
        proposed_state_parts,
        proposed_target_log_prob,
        proposed_target_log_prob_grad_parts,
    ]

Applies num_leapfrog_steps of the leapfrog integrator .


def _compute_log_acceptance_correction(current_momentums,
                                       proposed_momentums,
                                       independent_chain_ndims,
                                       name=None):
  """Helper to `kernel` which computes the log acceptance-correction.

  A sufficient but not necessary condition for the existence of a stationary
  distribution, `p(x)`, is "detailed balance", i.e.:

  ```none
  p(x'|x) p(x) = p(x|x') p(x')
  ```

  In the Metropolis-Hastings algorithm, a state is proposed according to
  `g(x'|x)` and accepted according to `a(x'|x)`, hence
  `p(x'|x) = g(x'|x) a(x'|x)`.

  Inserting this into the detailed balance equation implies:

  ```none
      g(x'|x) a(x'|x) p(x) = g(x|x') a(x|x') p(x')
  ==> a(x'|x) / a(x|x') = p(x') / p(x) [g(x|x') / g(x'|x)]    (*)
  ```

  One definition of `a(x'|x)` which satisfies (*) is:

  ```none
  a(x'|x) = min(1, p(x') / p(x) [g(x|x') / g(x'|x)])
  ```

  (To see that this satisfies (*), notice that under this definition only at
  most one `a(x'|x)` and `a(x|x') can be other than one.)

  We call the bracketed term the "acceptance correction".

  In the case of UncalibratedHMC, the log acceptance-correction is not the log
  proposal-ratio. UncalibratedHMC augments the state-space with momentum, z.
  Assuming a standard Gaussian distribution for momentums, the chain eventually
  converges to:

  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  ```

  Relating this back to Metropolis-Hastings parlance, for HMC we have:

  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  g([x, z] | [x', z']) = g([x', z'] | [x, z])
  ```

  In other words, the MH bracketed term is `1`. However, because we desire to
  use a general MH framework, we can place the momentum probability ratio inside
  the metropolis-correction factor thus getting an acceptance probability:

  ```none
                       target_prob(x')
  accept_prob(x'|x) = -----------------  [exp(-0.5 z**2) / exp(-0.5 z'**2)]
                       target_prob(x)
  ```

  (Note: we actually need to handle the kinetic energy change at each leapfrog
  step, but this is the idea.)

  Args:
    current_momentums: `Tensor` representing the value(s) of the current
      momentum(s) of the state (parts).
    proposed_momentums: `Tensor` representing the value(s) of the proposed
      momentum(s) of the state (parts).
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'compute_log_acceptance_correction').

  Returns:
    log_acceptance_correction: `Tensor` representing the `log`
      acceptance-correction.  (See docstring for mathematical definition.)
  """
  with tf.compat.v1.name_scope(
      name, 'compute_log_acceptance_correction',
      [independent_chain_ndims, current_momentums, proposed_momentums]):
    log_current_kinetic, log_proposed_kinetic = [], []
    for current_momentum, proposed_momentum in zip(
        current_momentums, proposed_momentums):
      axis = tf.range(independent_chain_ndims, tf.rank(current_momentum))
      log_current_kinetic.append(_log_sum_sq(current_momentum, axis))
      log_proposed_kinetic.append(_log_sum_sq(proposed_momentum, axis))
    current_kinetic = 0.5 * tf.exp(
        tf.reduce_logsumexp(
            input_tensor=tf.stack(log_current_kinetic, axis=-1), axis=-1))
    proposed_kinetic = 0.5 * tf.exp(
        tf.reduce_logsumexp(
            input_tensor=tf.stack(log_proposed_kinetic, axis=-1), axis=-1))
    return mcmc_util.safe_sum([current_kinetic, -proposed_kinetic])

def _compute_log_acceptance_correction ( current_momentums , proposed_momentums , independent_chain_ndims , name = None ) : with tf . compat . v1 . name_scope ( name , 'compute_log_acceptance_correction' , [ independent_chain_ndims , current_momentums , proposed_momentums ] ) : log_current_kinetic , log_proposed_kinetic = [ ] , [ ] for current_momentum , proposed_momentum in zip ( current_momentums , proposed_momentums ) : axis = tf . range ( independent_chain_ndims , tf . rank ( current_momentum ) ) log_current_kinetic . append ( _log_sum_sq ( current_momentum , axis ) ) log_proposed_kinetic . append ( _log_sum_sq ( proposed_momentum , axis ) ) current_kinetic = 0.5 * tf . exp ( tf . reduce_logsumexp ( input_tensor = tf . stack ( log_current_kinetic , axis = - 1 ) , axis = - 1 ) ) proposed_kinetic = 0.5 * tf . exp ( tf . reduce_logsumexp ( input_tensor = tf . stack ( log_proposed_kinetic , axis = - 1 ) , axis = - 1 ) ) return mcmc_util . safe_sum ( [ current_kinetic , - proposed_kinetic ] )

Helper to kernel which computes the log acceptance - correction .


def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of Hamiltonian Monte Carlo.

    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`.
    """
    previous_step_size_assign = (
        [] if self.step_size_update_fn is None
        else (previous_kernel_results.extra.step_size_assign
              if mcmc_util.is_list_like(
                  previous_kernel_results.extra.step_size_assign)
              else [previous_kernel_results.extra.step_size_assign]))

    with tf.control_dependencies(previous_step_size_assign):
      next_state, kernel_results = self._impl.one_step(
          current_state, previous_kernel_results)
      if self.step_size_update_fn is not None:
        step_size_assign = self.step_size_update_fn(  # pylint: disable=not-callable
            self.step_size, kernel_results)
        kernel_results = kernel_results._replace(
            extra=HamiltonianMonteCarloExtraKernelResults(
                step_size_assign=step_size_assign))
      return next_state, kernel_results

def one_step ( self , current_state , previous_kernel_results ) : previous_step_size_assign = ( [ ] if self . step_size_update_fn is None else ( previous_kernel_results . extra . step_size_assign if mcmc_util . is_list_like ( previous_kernel_results . extra . step_size_assign ) else [ previous_kernel_results . extra . step_size_assign ] ) ) with tf . control_dependencies ( previous_step_size_assign ) : next_state , kernel_results = self . _impl . one_step ( current_state , previous_kernel_results ) if self . step_size_update_fn is not None : step_size_assign = self . step_size_update_fn ( self . step_size , kernel_results ) kernel_results = kernel_results . _replace ( extra = HamiltonianMonteCarloExtraKernelResults ( step_size_assign = step_size_assign ) ) return next_state , kernel_results

Runs one iteration of Hamiltonian Monte Carlo .


def bootstrap_results(self, init_state):
    """Creates initial `previous_kernel_results` using a supplied `state`."""
    kernel_results = self._impl.bootstrap_results(init_state)
    if self.step_size_update_fn is not None:
      step_size_assign = self.step_size_update_fn(self.step_size, None)  # pylint: disable=not-callable
      kernel_results = kernel_results._replace(
          extra=HamiltonianMonteCarloExtraKernelResults(
              step_size_assign=step_size_assign))
    return kernel_results

def bootstrap_results ( self , init_state ) : kernel_results = self . _impl . bootstrap_results ( init_state ) if self . step_size_update_fn is not None : step_size_assign = self . step_size_update_fn ( self . step_size , None ) kernel_results = kernel_results . _replace ( extra = HamiltonianMonteCarloExtraKernelResults ( step_size_assign = step_size_assign ) ) return kernel_results

Creates initial previous_kernel_results using a supplied state .


def bayesian_resnet(input_shape,
                    num_classes=10,
                    kernel_posterior_scale_mean=-9.0,
                    kernel_posterior_scale_stddev=0.1,
                    kernel_posterior_scale_constraint=0.2):
  """Constructs a ResNet18 model.

  Args:
    input_shape: A `tuple` indicating the Tensor shape.
    num_classes: `int` representing the number of class labels.
    kernel_posterior_scale_mean: Python `int` number for the kernel
      posterior's scale (log variance) mean. The smaller the mean the closer
      is the initialization to a deterministic network.
    kernel_posterior_scale_stddev: Python `float` number for the initial kernel
      posterior's scale stddev.
      ```
      q(W|x) ~ N(mu, var),
      log_var ~ N(kernel_posterior_scale_mean, kernel_posterior_scale_stddev)
      ````
    kernel_posterior_scale_constraint: Python `float` number for the log value
      to constrain the log variance throughout training.
      i.e. log_var <= log(kernel_posterior_scale_constraint).

  Returns:
    tf.keras.Model.
  """

  filters = [64, 128, 256, 512]
  kernels = [3, 3, 3, 3]
  strides = [1, 2, 2, 2]

  def _untransformed_scale_constraint(t):
    return tf.clip_by_value(t, -1000,
                            tf.math.log(kernel_posterior_scale_constraint))

  kernel_posterior_fn = tfp.layers.default_mean_field_normal_fn(
      untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(
          mean=kernel_posterior_scale_mean,
          stddev=kernel_posterior_scale_stddev),
      untransformed_scale_constraint=_untransformed_scale_constraint)

  image = tf.keras.layers.Input(shape=input_shape, dtype='float32')
  x = tfp.layers.Convolution2DFlipout(
      64,
      3,
      strides=1,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(image)

  for i in range(len(kernels)):
    x = _resnet_block(
        x,
        filters[i],
        kernels[i],
        strides[i],
        kernel_posterior_fn)

  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)
  x = tf.keras.layers.AveragePooling2D(4, 1)(x)
  x = tf.keras.layers.Flatten()(x)

  x = tfp.layers.DenseFlipout(
      num_classes,
      kernel_posterior_fn=kernel_posterior_fn)(x)

  model = tf.keras.Model(inputs=image, outputs=x, name='resnet18')
  return model

def bayesian_resnet ( input_shape , num_classes = 10 , kernel_posterior_scale_mean = - 9.0 , kernel_posterior_scale_stddev = 0.1 , kernel_posterior_scale_constraint = 0.2 ) : filters = [ 64 , 128 , 256 , 512 ] kernels = [ 3 , 3 , 3 , 3 ] strides = [ 1 , 2 , 2 , 2 ] def _untransformed_scale_constraint ( t ) : return tf . clip_by_value ( t , - 1000 , tf . math . log ( kernel_posterior_scale_constraint ) ) kernel_posterior_fn = tfp . layers . default_mean_field_normal_fn ( untransformed_scale_initializer = tf . compat . v1 . initializers . random_normal ( mean = kernel_posterior_scale_mean , stddev = kernel_posterior_scale_stddev ) , untransformed_scale_constraint = _untransformed_scale_constraint ) image = tf . keras . layers . Input ( shape = input_shape , dtype = 'float32' ) x = tfp . layers . Convolution2DFlipout ( 64 , 3 , strides = 1 , padding = 'same' , kernel_posterior_fn = kernel_posterior_fn ) ( image ) for i in range ( len ( kernels ) ) : x = _resnet_block ( x , filters [ i ] , kernels [ i ] , strides [ i ] , kernel_posterior_fn ) x = tf . keras . layers . BatchNormalization ( ) ( x ) x = tf . keras . layers . Activation ( 'relu' ) ( x ) x = tf . keras . layers . AveragePooling2D ( 4 , 1 ) ( x ) x = tf . keras . layers . Flatten ( ) ( x ) x = tfp . layers . DenseFlipout ( num_classes , kernel_posterior_fn = kernel_posterior_fn ) ( x ) model = tf . keras . Model ( inputs = image , outputs = x , name = 'resnet18' ) return model