class_docstr: Sequential Monte Carlo transition kernel.
  
  Sequential Monte Carlo maintains a population of weighted particles
  representing samples from a sequence of target distributions. It is
  *not* a calibrated MCMC kernel: the transitions step through a sequence of
  target distributions, rather than trying to maintain a stationary
  distribution.
class_docstr_tok: ["Sequential","Monte","Carlo","transition","kernel",".","Sequential","Monte","Carlo","maintains","a","population","of","weighted","particles","representing","samples","from","a","sequence","of","target","distributions",".","It","is","*","not","*","a","calibrated","MCMC","kernel",":","the","transitions","step","through","a","sequence","of","target","distributions",",","rather","than","trying","to","maintain","a","stationary","distribution","."]
docstr: Takes one Sequential Monte Carlo inference step.
  
  Args:
    state: instance of `tfp.experimental.mcmc.WeightedParticles` representing
      the current particles with (log) weights. The `log_weights` must be
      a float `Tensor` of shape `[num_particles, b1, ..., bN]`. The
      `particles` may be any structure of `Tensor`s, each of which
      must have shape `concat([log_weights.shape, event_shape])` for some
      `event_shape`, which may vary across components.
    kernel_results: instance of
      `tfp.experimental.mcmc.SequentialMonteCarloResults` representing results
      from a previous step.
    seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
  Returns:
    state: instance of `tfp.experimental.mcmc.WeightedParticles` representing
      new particles with (log) weights.
    kernel_results: instance of
      `tfp.experimental.mcmc.SequentialMonteCarloResults`.
docstr_tok: ["Takes","one","Sequential","Monte","Carlo","inference","step",".","Args",":","state",":","instance","of","`","tfp.experimental.mcmc.WeightedParticles","`","representing","the","current","particles","with","(","log",")","weights",".","The","`","log_weights","`","must","be","a","float","`","Tensor","`","of","shape","`","[","num_particles",",","b1",",","...",",","bN","]","`",".","The","`","particles","`","may","be","any","structure","of","`","Tensor","`","s",",","each","of","which","must","have","shape","`","concat","(","[","log_weights.shape",",","event_shape","]",")","`","for","some","`","event_shape","`",",","which","may","vary","across","components",".","kernel_results",":","instance","of","`","tfp.experimental.mcmc.SequentialMonteCarloResults","`","representing","results","from","a","previous","step",".","seed",":","PRNG","seed",";","see","`","tfp.random.sanitize_seed","`","for","details",".","Returns",":","state",":","instance","of","`","tfp.experimental.mcmc.WeightedParticles","`","representing","new","particles","with","(","log",")","weights",".","kernel_results",":","instance","of","`","tfp.experimental.mcmc.SequentialMonteCarloResults","`","."]
code: def one_step(self, state, kernel_results, seed=None):
      """Takes one Sequential Monte Carlo inference step.
  
      Args:
        state: instance of `tfp.experimental.mcmc.WeightedParticles` representing
          the current particles with (log) weights. The `log_weights` must be
          a float `Tensor` of shape `[num_particles, b1, ..., bN]`. The
          `particles` may be any structure of `Tensor`s, each of which
          must have shape `concat([log_weights.shape, event_shape])` for some
          `event_shape`, which may vary across components.
        kernel_results: instance of
          `tfp.experimental.mcmc.SequentialMonteCarloResults` representing results
          from a previous step.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
      Returns:
        state: instance of `tfp.experimental.mcmc.WeightedParticles` representing
          new particles with (log) weights.
        kernel_results: instance of
          `tfp.experimental.mcmc.SequentialMonteCarloResults`.
      """
      with tf.name_scope(self.name):
          with tf.name_scope("one_step"):
              seed = samplers.sanitize_seed(seed)
              proposal_seed, resample_seed = samplers.split_seed(seed)
  
              state = WeightedParticles(*state)  # Canonicalize.
  
              # Propose new particles and update weights for this step, unless it's
              # the initial step, in which case, use the user-provided initial
              # particles and weights.
              proposed_state = self.propose_and_update_log_weights_fn(
                  # Propose state[t] from state[t - 1].
                  ps.maximum(0, kernel_results.steps - 1),
                  state,
                  seed=proposal_seed,
              )
              is_initial_step = ps.equal(kernel_results.steps, 0)
              # TODO(davmre): this `where` assumes the state size didn't change.
              state = tf.nest.map_structure(
                  lambda a, b: tf.where(is_initial_step, a, b),
                  state,
                  proposed_state,
              )
  
              normalized_log_weights = tf.nn.log_softmax(
                  state.log_weights, axis=self.particles_dim
              )
              # Every entry of `log_weights` differs from `normalized_log_weights`
              # by the same normalizing constant. We extract that constant by
              # examining an arbitrary entry.
              incremental_log_marginal_likelihood = tf.gather(
                  state.log_weights, 0, axis=self.particles_dim
              ) - tf.gather(
                  normalized_log_weights, 0, axis=self.particles_dim
              )
  
              do_resample = self.resample_criterion_fn(
                  state, particles_dim=self.particles_dim
              )
              # Some batch elements may require resampling and others not, so
              # we first do the resampling for all elements, then select whether to
              # use the resampled values for each batch element according to
              # `do_resample`. If there were no batching, we might prefer to use
              # `tf.cond` to avoid the resampling computation on steps where it's not
              # needed---but we're ultimately interested in adaptive resampling
              # for statistical (not computational) purposes, so this isn't a
              # dealbreaker.
              [
                  resampled_particles,
                  resample_indices,
                  weights_after_resampling,
              ] = weighted_resampling.resample(
                  particles=state.particles,
                  # The `stop_gradient` here does not affect discrete resampling
                  # (which is nondifferentiable anyway), but avoids canceling out
                  # the gradient signal from the 'target' log weights, as described in
                  # Scibior, Masrani, and Wood (2021).
                  log_weights=tf.stop_gradient(state.log_weights),
                  resample_fn=self.resample_fn,
                  target_log_weights=(
                      normalized_log_weights
                      if self.unbiased_gradients
                      else None
                  ),
                  particles_dim=self.particles_dim,
                  seed=resample_seed,
              )
              (
                  resampled_particles,
                  resample_indices,
                  log_weights,
              ) = tf.nest.map_structure(
                  lambda r, p: mcmc_util.choose(do_resample, r, p),
                  (
                      resampled_particles,
                      resample_indices,
                      weights_after_resampling,
                  ),
                  (
                      state.particles,
                      _dummy_indices_like(resample_indices),
                      normalized_log_weights,
                  ),
              )
  
          return (
              WeightedParticles(
                  particles=resampled_particles, log_weights=log_weights
              ),
              SequentialMonteCarloResults(
                  steps=kernel_results.steps + 1,
                  parent_indices=resample_indices,
                  incremental_log_marginal_likelihood=(
                      incremental_log_marginal_likelihood
                  ),
                  accumulated_log_marginal_likelihood=(
                      kernel_results.accumulated_log_marginal_likelihood
                      + incremental_log_marginal_likelihood
                  ),
                  seed=seed,
              ),
          )
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import weighted_resampling
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 9 for calling the SequentialMonteCarlo obj's underlying member method code with example usage: obj.one_step(state, kernel_results, seed) and returns: WeightedParticles,  SequentialMonteCarloResults