code: def sequential_monte_carlo(
      initial_weighted_particles,
      propose_and_update_log_weights_fn,
      num_steps,
      resample_fn,
      resample_criterion_fn,
      trace_fn,
      trace_criterion_fn,
      particles_dim=0,
      parallel_iterations=1,
      unbiased_gradients=True,
      static_trace_allocation_size=None,
      seed=None,
      name=None,
  ):
      """Run Sequential Monte Carlo.
  
      Sequential Monte Carlo maintains a population of weighted particles
      representing samples from a sequence of target distributions.
  
      Args:
        initial_weighted_particles: The initial
          `tfp.experimental.mcmc.WeightedParticles`.
        propose_and_update_log_weights_fn: Python `callable` with signature
          `new_weighted_particles = propose_and_update_log_weights_fn(step,
          weighted_particles, seed=None)`. Its input is a
          `tfp.experimental.mcmc.WeightedParticles` structure representing
          weighted samples (with normalized weights) from the `step`th
          target distribution, and it returns another such structure representing
          unnormalized weighted samples from the next (`step + 1`th) target
          distribution. This will typically include particles
          sampled from a proposal distribution `q(x[step + 1] | x[step])`, and
          weights that account for some or all of: the proposal density,
          a transition density `p(x[step + 1] | x[step]),
          observation weights `p(y[step + 1] | x[step + 1])`, and/or a backwards
          or 'L'-kernel `L(x[step] | x[step + 1])`. The (log) normalization
          constant of the weights is interpreted as the incremental (log) marginal
          likelihood.
        num_steps: Number of steps to run Sequential Monte Carlo.
        resample_fn: Resampling scheme specified as a `callable` with signature
          `indices = resample_fn(log_probs, event_size, sample_shape, seed)`,
          where `log_probs` is a `Tensor` of the same shape as `state.log_weights`
          containing a normalized log-probability for every current
          particle, `event_size` is the number of new particle indices to
          generate,  `sample_shape` is the number of independent index sets to
          return, and the  return value `indices` is an `int` Tensor of shape
          `concat([sample_shape, [event_size, B1, ..., BN])`. Typically one of
          `tfp.experimental.mcmc.resample_deterministic_minimum_error`,
          `tfp.experimental.mcmc.resample_independent`,
          `tfp.experimental.mcmc.resample_stratified`, or
          `tfp.experimental.mcmc.resample_systematic`.
          Default value: `tfp.experimental.mcmc.resample_systematic`.
        resample_criterion_fn: optional Python `callable` with signature
          `do_resample = resample_criterion_fn(weighted_particles)`,
          passed an instance of `tfp.experimental.mcmc.WeightedParticles`. The
          return value `do_resample`
          determines whether particles are resampled at the current step. The
          default behavior is to resample particles when the effective
          sample size falls below half of the total number of particles.
          Default value: `tfp.experimental.mcmc.ess_below_threshold`.
        trace_fn: Python `callable` defining the values to be traced at each step,
          with signature `traced_values = trace_fn(weighted_particles, results)`
          in which the first argument is an instance of
          `tfp.experimental.mcmc.WeightedParticles` and the second an instance of
          `SequentialMonteCarloResults` tuple, and the return value is a structure
          of `Tensor`s.
          Default value: `lambda s, r: (s.particles, s.log_weights,
          r.parent_indices, r.incremental_log_marginal_likelihood)`
        trace_criterion_fn: optional Python `callable` with signature
          `trace_this_step = trace_criterion_fn(weighted_particles, results)`
          taking the same arguments as `trace_fn` and returning a boolean `Tensor`.
          If `None`, only values from the final step are returned.
          Default value: `lambda *_: True` (trace every step).
        particles_dim: `int` dimension that indexes the particles in the
          `tfp.experimental.mcmc.WeightedParticles` structures on which this
          function operates.
          Default value: `0`.
        parallel_iterations: Passed to the internal `tf.while_loop`.
          Default value: `1`.
        unbiased_gradients: If `True`, use the stop-gradient
          resampling trick of Scibior, Masrani, and Wood [1] to correct for
          gradient bias introduced by the discrete resampling step. This will
          generally increase the variance of stochastic gradients.
          Default value: `True`.
        static_trace_allocation_size: Optional Python `int` size of trace to
          allocate statically. This should be an upper bound on the number of steps
          traced and is used only when the length cannot be
          statically inferred (for example, if a `trace_criterion_fn` is
          specified).
          It is primarily intended for contexts where static shapes are required,
          such as in XLA-compiled code.
          Default value: `None`.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        name: Python `str` name for ops created by this method.
          Default value: `None` (i.e., `'particle_filter'`).
      Returns:
        traced_results: A structure of Tensors as returned by `trace_fn`. If
          `trace_criterion_fn==None`, this is computed from the final step;
          otherwise, each Tensor will have initial dimension `num_steps_traced`
          and stacks the traced results across all steps.
  
      #### References
  
      [1] Adam Scibior, Vaden Masrani, and Frank Wood. Differentiable Particle
          Filtering without Modifying the Forward Pass. _arXiv preprint
          arXiv:2106.10314_, 2021. https://arxiv.org/abs/2106.10314
      """
      with tf.name_scope(name or "sequential_monte_carlo"):
          seed = samplers.sanitize_seed(seed)
  
          kernel = smc_kernel.SequentialMonteCarlo(
              propose_and_update_log_weights_fn=propose_and_update_log_weights_fn,
              resample_fn=resample_fn,
              resample_criterion_fn=resample_criterion_fn,
              particles_dim=particles_dim,
              unbiased_gradients=unbiased_gradients,
          )
  
          # If trace criterion is `None`, we'll return only the final results.
          never_trace = lambda *_: False
          if trace_criterion_fn is None:
              static_trace_allocation_size = 0
              trace_criterion_fn = never_trace
  
          # Use `trace_scan` rather than `sample_chain` directly because the latter
          # would force us to trace the state history (with or without thinning),
          # which is not always appropriate.
          def seeded_one_step(seed_state_results, _):
              seed, state, results = seed_state_results
              one_step_seed, next_seed = samplers.split_seed(seed)
              next_state, next_results = kernel.one_step(
                  state, results, seed=one_step_seed
              )
              return next_seed, next_state, next_results
  
          (
              final_seed_state_result,
              traced_results,
          ) = loop_util.trace_scan(
              loop_fn=seeded_one_step,
              initial_state=(
                  seed,
                  initial_weighted_particles,
                  kernel.bootstrap_results(initial_weighted_particles),
              ),
              elems=tf.ones([num_steps]),
              trace_fn=lambda seed_state_results: trace_fn(
                  *seed_state_results[1:]
              ),
              trace_criterion_fn=(
                  lambda seed_state_results: trace_criterion_fn(  # pylint: disable=g-long-lambda
                      *seed_state_results[1:]
                  )
              ),
              static_trace_allocation_size=static_trace_allocation_size,
              parallel_iterations=parallel_iterations,
          )
  
          if trace_criterion_fn is never_trace:
              # Return results from just the final step.
              traced_results = trace_fn(*final_seed_state_result[1:])
  
          return traced_results
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/particle_filter.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import sequential_monte_carlo_kernel
  from tensorflow_probability.python.experimental.mcmc import weighted_resampling
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import docstring_util
  from tensorflow_probability.python.internal import loop_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 3 for calling a global function with example usage: sequential_monte_carlo(initial_weighted_particles, propose_and_update_log_weights_fn, num_steps, resample_fn, resample_criterion_fn, trace_fn, trace_criterion_fn, particles_dim, parallel_iterations, unbiased_gradients, static_trace_allocation_size, seed, name) and returns: traced_results, next_seed,  next_state,  next_results