code: def _sample_next(
      target_log_prob_fn,
      current_state_parts,
      step_sizes,
      max_doublings,
      current_target_log_prob,
      batch_rank,
      seed=None,
      experimental_shard_axis_names=None,
      name=None,
  ):
      """Applies a single iteration of slice sampling update.
  
      Applies hit and run style slice sampling. Chooses a uniform random direction
      on the unit sphere in the event space. Applies the one dimensional slice
      sampling update along that direction.
  
      Args:
        target_log_prob_fn: Python callable which takes an argument like
          `*current_state_parts` and returns its (possibly unnormalized) log-density
          under the target distribution.
        current_state_parts: Python `list` of `Tensor`s representing the current
          state(s) of the Markov chain(s). The first `independent_chain_ndims` of
          the `Tensor`(s) index different chains.
        step_sizes: Python `list` of `Tensor`s. Provides a measure of the width
          of the density. Used to find the slice bounds. Must broadcast with the
          shape of `current_state_parts`.
        max_doublings: Integer number of doublings to allow while locating the slice
          boundaries.
        current_target_log_prob: `Tensor` representing the value of
          `target_log_prob_fn(*current_state_parts)`. The only reason to specify
          this argument is to reduce TF graph size.
        batch_rank: Integer. The number of axes in the state that correspond to
          independent batches.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        experimental_shard_axis_names: A structure of string names indicating how
          members of the state are sharded.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., 'find_slice_bounds').
  
      Returns:
        proposed_state_parts: Tensor or Python list of `Tensor`s representing the
          state(s) of the Markov chain(s) at each result step. Has same shape as
          input `current_state_parts`.
        proposed_target_log_prob: `Tensor` representing the value of
          `target_log_prob_fn` at `next_state`.
        bounds_satisfied: Boolean `Tensor` of the same shape as the log density.
          True indicates whether the an interval containing the slice for that
          batch was found successfully.
        direction: `Tensor` or Python list of `Tensors`s representing the direction
          along which the slice was sampled. Has the same shape and dtype(s) as
          `current_state_parts`.
        upper_bounds: `Tensor` of batch shape and the dtype of the input state. The
          upper bounds of the slices along the sampling direction.
        lower_bounds: `Tensor` of batch shape and the dtype of the input state. The
          lower bounds of the slices along the sampling direction.
      """
      direction_seed, slice_seed = samplers.split_seed(seed)
      with tf.name_scope(name or "sample_next"):
          # First step: Choose a random direction.
          # Direction is a list of tensors. The i'th tensor should have the same shape
          # as the i'th state part.
          direction = _choose_random_direction(
              current_state_parts,
              batch_rank=batch_rank,
              seed=direction_seed,
              experimental_shard_axis_names=experimental_shard_axis_names,
          )
  
          # Interpolates the step sizes for the chosen direction.
          # Applies an ellipsoidal interpolation to compute the step direction for
          # the chosen direction. Suppose we are given step sizes for each direction.
          # Label these s_1, s_2, ... s_k. These are the step sizes to use if moving
          # in a direction parallel to one of the axes. Consider an ellipsoid which
          # intercepts the i'th axis at s_i. The step size for a direction specified
          # by the unit vector (n_1, n_2 ...n_k) is then defined as the intersection
          # of the line through this vector with this ellipsoid.
          #
          # One can show that the length of the vector from the origin to the
          # intersection point is given by:
          # 1 / sqrt(n_1^2 / s_1^2  + n_2^2 / s_2^2  + ...).
          #
          # Proof:
          # The equation of the ellipsoid is:
          # Sum_i [x_i^2 / s_i^2 ] = 1. Let n be a unit direction vector. Points
          # along the line given by n may be parameterized as alpha*n where alpha is
          # the distance along the vector. Plugging this into the equation for the
          # ellipsoid, we get:
          # alpha^2 ( n_1^2 / s_1^2 + n_2^2 / s_2^2 + ...) = 1
          # so alpha = \sqrt { \frac{1} { ( n_1^2 / s_1^2 + n_2^2 / s_2^2 + ...) } }
          reduce_axes = [
              ps.range(batch_rank, ps.rank(dirn_part))
              for dirn_part in direction
          ]
          experimental_shard_axis_names = (
              experimental_shard_axis_names or ([None] * len(direction))
          )
  
          def reduce_sum(v, axis, shard_axes):
              out = tf.reduce_sum(v, axis=axis)
              if shard_axes is not None:
                  out = distribute_lib.psum(out, shard_axes)
              return out
  
          components = [
              reduce_sum(
                  (dirn_part / step_size) ** 2,
                  reduce_axes[i],
                  shard_axes,
              )
              for i, (step_size, dirn_part, shard_axes) in enumerate(
                  zip(
                      step_sizes,
                      direction,
                      experimental_shard_axis_names,
                  )
              )
          ]
          step_size = tf.math.rsqrt(tf.add_n(components))
          # Computes the rank of a tensor. Uses the static rank if possible.
          state_part_ranks = [
              ps.rank(part) for part in current_state_parts
          ]
  
          def _step_along_direction(alpha):
              """Converts the scalar alpha into an n-dim vector with full state info.
  
              Computes x_0 + alpha * direction where x_0 is the current state and
              direction is the direction chosen above.
  
              Args:
                alpha: A tensor of shape equal to the batch dimensions of
                  `current_state_parts`.
  
              Returns:
                state_parts: Tensor or Python list of `Tensor`s representing the
                  state(s) of the Markov chain(s) for a given alpha and a given chosen
                  direction. Has the same shape as `current_state_parts`.
              """
              padded_alphas = [
                  _right_pad(alpha, final_rank=part_rank)
                  for part_rank in state_part_ranks
              ]
  
              state_parts = [
                  state_part + padded_alpha * direction_part
                  for state_part, direction_part, padded_alpha in zip(
                      current_state_parts, direction, padded_alphas
                  )
              ]
              return state_parts
  
          def projected_target_log_prob_fn(alpha):
              """The target log density projected along the chosen direction.
  
              Args:
                alpha: A tensor of shape equal to the batch dimensions of
                  `current_state_parts`.
  
              Returns:
                Target log density evaluated at x_0 + alpha * direction where x_0 is the
                current state and direction is the direction chosen above. Has the same
                shape as `alpha`.
              """
              return target_log_prob_fn(*_step_along_direction(alpha))
  
          alpha_init = tf.zeros_like(
              current_target_log_prob,
              dtype=current_state_parts[0].dtype,
          )
          [
              next_alpha,
              next_target_log_prob,
              bounds_satisfied,
              upper_bounds,
              lower_bounds,
          ] = ssu.slice_sampler_one_dim(
              projected_target_log_prob_fn,
              x_initial=alpha_init,
              max_doublings=max_doublings,
              step_size=step_size,
              seed=slice_seed,
          )
          return [
              _step_along_direction(next_alpha),
              next_target_log_prob,
              bounds_satisfied,
              direction,
              upper_bounds,
              lower_bounds,
          ]
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import slice_sampler_utils
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow_probability.python.random import random_ops
label_desc: Use image node_id 2 for calling a global function with example usage: _sample_next(target_log_prob_fn, current_state_parts, step_sizes, max_doublings, current_target_log_prob, batch_rank, seed, experimental_shard_axis_names, name) and returns: list, out, state_parts, target_log_prob_fn