code: def _particle_filter_propose_and_update_log_weights_fn(
      observations,
      transition_fn,
      proposal_fn,
      observation_fn,
      num_transitions_per_observation=1,
      particles_dim=0,
  ):
      """Build a function specifying a particle filter update step."""
  
      def propose_and_update_log_weights_fn(step, state, seed=None):
          particles, log_weights = state.particles, state.log_weights
          transition_dist = transition_fn(step, particles)
          assertions = _assert_batch_shape_matches_weights(
              distribution=transition_dist,
              weights_shape=ps.shape(log_weights),
              diststr="transition",
          )
  
          if proposal_fn:
              proposal_dist = proposal_fn(step, particles)
              assertions += _assert_batch_shape_matches_weights(
                  distribution=proposal_dist,
                  weights_shape=ps.shape(log_weights),
                  diststr="proposal",
              )
              proposed_particles = proposal_dist.sample(seed=seed)
  
              log_weights += transition_dist.log_prob(
                  proposed_particles
              ) - proposal_dist.log_prob(proposed_particles)
              # The normalizing constant E~q[p(x)/q(x)] is 1 in expectation,
              # so we reduce variance by dividing it out. Intuitively: the marginal
              # likelihood of a model with no observations is constant
              # (equal to 1.), so the transition and proposal distributions shouldn't
              # affect it.
              log_weights = tf.nn.log_softmax(
                  log_weights, axis=particles_dim
              )
          else:
              proposed_particles = transition_dist.sample(seed=seed)
  
          with tf.control_dependencies(assertions):
              return smc_kernel.WeightedParticles(
                  particles=proposed_particles,
                  log_weights=log_weights
                  + _compute_observation_log_weights(
                      step + 1,
                      proposed_particles,
                      observations,
                      observation_fn,
                      num_transitions_per_observation=num_transitions_per_observation,
                      particles_dim=particles_dim,
                  ),
              )
  
      return propose_and_update_log_weights_fn
  
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 6 for calling a global function with example usage: _particle_filter_propose_and_update_log_weights_fn(observations, transition_fn, proposal_fn, observation_fn, num_transitions_per_observation, particles_dim) and returns: propose_and_update_log_weights_fn, smc_kernel