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

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

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

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

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

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

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

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

  Returns:
    x_proposed: A tensor of the same shape and dtype as `x_initial`. The next
      proposed state of the chain.
Samples from the slice by applying shrinkage for rejected points .