code: def resample_stratified(
      log_probs, event_size, sample_shape, seed=None, name=None
  ):
      """Stratified resampler for sequential Monte Carlo.
  
      The value returned from this algorithm is similar to sampling with
      ```python
      expanded_sample_shape = tf.concat([[event_size], sample_shape]), axis=-1)
      tfd.Categorical(logits=log_probs).sample(expanded_sample_shape)`
      ```
      but with values sorted along the first axis. It can be considered to be
      sampling events made up of a length-`event_size` vector of draws from
      the `Categorical` distribution. However, although the elements of
      this event have the appropriate marginal distribution, they are not
      independent of each other. Instead they are drawn using a low variance
      stratified sampling method suitable for use with Sequential Monte
      Carlo algorithms.
      The sortedness is an unintended side effect of the algorithm that is
      harmless in the context of simple SMC algorithms.
  
      This function is based on Algorithm #1 in the paper
      [Maskell et al. (2006)][1].
  
      Args:
        log_probs: A tensor-valued batch of discrete log probability distributions.
        event_size: the dimension of the vector considered a single draw.
        sample_shape: the `sample_shape` determining the number of draws.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
          Default value: None (i.e. no seed).
        name: Python `str` name for ops created by this method.
          Default value: `None` (i.e., `'resample_independent'`).
  
      Returns:
        resampled_indices: a tensor of samples.
  
      #### References
  
      [1]: S. Maskell, B. Alun-Jones and M. Macleod. A Single Instruction Multiple
           Data Particle Filter.
           In 2006 IEEE Nonlinear Statistical Signal Processing Workshop.
           http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
  
      """
      with tf.name_scope(name or "resample_stratified") as name:
          log_probs = tf.convert_to_tensor(
              log_probs, dtype_hint=tf.float32
          )
          log_probs = dist_util.move_dimension(
              log_probs, source_idx=0, dest_idx=-1
          )
          points_shape = ps.concat(
              [sample_shape, ps.shape(log_probs)[:-1], [event_size]],
              axis=0,
          )
          # Draw an offset for every element of an event.
          interval_width = ps.cast(
              1.0 / event_size, dtype=log_probs.dtype
          )
          offsets = uniform.Uniform(
              low=ps.cast(0.0, dtype=log_probs.dtype),
              high=interval_width,
          ).sample(points_shape, seed=seed)
          # The unit interval is divided into equal partitions and each point
          # is a random offset into a partition.
          even_spacing = (
              tf.linspace(
                  start=ps.cast(0.0, dtype=log_probs.dtype),
                  stop=1 - interval_width,
                  num=event_size,
              )
              + offsets
          )
          log_points = tf.math.log(even_spacing)
  
          resampled = _resample_using_log_points(
              log_probs, sample_shape, log_points
          )
          return dist_util.move_dimension(
              resampled, source_idx=-1, dest_idx=0
          )
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/weighted_resampling.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.gradient import value_and_gradient
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 5 for calling a global function with example usage: resample_stratified(log_probs, event_size, sample_shape, seed, name) and returns: dist_util