code: def _sample_with_shrinkage(
      x_initial,
      target_log_prob,
      log_slice_heights,
      step_size,
      lower_bounds,
      upper_bounds,
      seed,
      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: PRNG seed; see `tfp.random.sanitize_seed` for details.
        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.name_scope(name or "sample_with_shrinkage"):
          seed = samplers.sanitize_seed(seed)
          # 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(found)
          x_next = tf.identity(x_initial)
          x_initial_shape = ps.shape(x_initial)
          x_initial_dtype = dtype_util.base_dtype(x_initial.dtype)
  
          def _body(found, seed, left, right, x_next):
              """Iterates until every chain has found a suitable next state."""
              proportions_seed, next_seed = samplers.split_seed(seed)
              proportions = samplers.uniform(
                  x_initial_shape,
                  dtype=x_initial_dtype,
                  seed=proportions_seed,
              )
              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_seed,
                  next_left,
                  next_right,
                  x_proposed,
              )
  
          return tf.while_loop(
              cond=cond,
              body=_body,
              loop_vars=(
                  found,
                  seed,
                  lower_bounds,
                  upper_bounds,
                  x_next,
              ),
          )[-1]
  
file_path: probability/tensorflow_probability/python/mcmc/internal/slice_sampler_utils.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
label_desc: Use image node_id 5 for calling a global function with example usage: _sample_with_shrinkage(x_initial, target_log_prob, log_slice_heights, step_size, lower_bounds, upper_bounds, seed, name) and returns: tf, next_found,  next_seed,  next_left,  next_right,  x_proposed