code: def resample(
      particles,
      log_weights,
      resample_fn,
      target_log_weights=None,
      particles_dim=0,
      seed=None,
  ):
      """Resamples the current particles according to provided weights.
  
      Args:
        particles: Nested structure of `Tensor`s each of shape
          `[num_particles, b1, ..., bN, ...]`, where
          `b1, ..., bN` are optional batch dimensions.
        log_weights: float `Tensor` of shape `[num_particles, b1, ..., bN]`, where
          `b1, ..., bN` are optional batch dimensions.
        resample_fn: choose the function used for resampling.
          Use 'resample_independent' for independent resamples.
          Use 'resample_stratified' for stratified resampling.
          Use 'resample_systematic' for systematic resampling.
        target_log_weights: optional float `Tensor` of the same shape and dtype as
          `log_weights`, specifying the target measure on `particles` if this is
          different from that implied by normalizing `log_weights`. The
          returned `log_weights_after_resampling` will represent this measure. If
          `None`, the target measure is implicitly taken to be the normalized
          log weights (`log_weights - tf.reduce_logsumexp(log_weights, axis=0)`).
          Default value: `None`.
        particles_dim: Python `int` axis of each state `Tensor` indexing into the
          particles. This is almost always zero, but nonzero values may be necessary
          when running SMC in nested contexts.
          Default value: `0`.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
      Returns:
        resampled_particles: Nested structure of `Tensor`s, matching `particles`.
        resample_indices: int `Tensor` of shape `[num_particles, b1, ..., bN]`.
        log_weights_after_resampling: float `Tensor` of same shape and dtype as
          `log_weights`, such that weighted sums of the resampled particles are
          equal (in expectation over the resampling step) to weighted sums of
          the original particles:
          `E [ exp(log_weights_after_resampling) * some_fn(resampled_particles) ]
          == exp(target_log_weights) * some_fn(particles)`.
          If no `target_log_weights` was specified, the log weights after
          resampling are uniformly equal to `-log(num_particles)`.
      """
      with tf.name_scope("resample"):
          num_particles = ps.dimension_size(log_weights, particles_dim)
  
          log_num_particles = tf.math.log(
              tf.cast(num_particles, log_weights.dtype)
          )
  
          # Normalize the weights and sample the ancestral indices.
          log_probs = tf.math.log_softmax(
              log_weights, axis=particles_dim
          )
          if particles_dim == 0:
              # For resample functions that don't yet support the
              # particles_dim argument.
              resampled_indices = resample_fn(
                  log_probs, num_particles, (), seed=seed
              )
          else:
              resampled_indices = resample_fn(
                  log_probs,
                  num_particles,
                  (),
                  particles_dim=particles_dim,
                  seed=seed,
              )
  
          gather_ancestors = (
              lambda x: (  # pylint: disable=g-long-lambda
                  mcmc_util.index_remapping_gather(
                      x,
                      resampled_indices,
                      axis=particles_dim,
                      indices_axis=particles_dim,
                  )
              )
          )
          resampled_particles = tf.nest.map_structure(
              gather_ancestors, particles
          )
          if target_log_weights is None:
              log_weights_after_resampling = tf.fill(
                  ps.shape(log_weights), -log_num_particles
              )
          else:
              importance_weights = (
                  target_log_weights - log_probs - log_num_particles
              )
              log_weights_after_resampling = tf.nest.map_structure(
                  gather_ancestors, importance_weights
              )
      return (
          resampled_particles,
          resampled_indices,
          log_weights_after_resampling,
      )
  
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 1 for calling a global function with example usage: resample(particles, log_weights, resample_fn, target_log_weights, particles_dim, seed) and returns: resampled_particles,  resampled_indices,  log_weights_after_resampling