code: def infer_trajectories(
      observations,
      initial_state_prior,
      transition_fn,
      observation_fn,
      num_particles,
      initial_state_proposal=None,
      proposal_fn=None,
      resample_fn=weighted_resampling.resample_systematic,
      resample_criterion_fn=smc_kernel.ess_below_threshold,
      unbiased_gradients=True,
      num_transitions_per_observation=1,
      particles_dim=0,
      seed=None,
      name=None,
  ):  # pylint: disable=g-doc-args
      """Use particle filtering to sample from the posterior over trajectories.
  
      ${particle_filter_arg_str}
        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., `'infer_trajectories'`).
      Returns:
        trajectories: a (structure of) Tensor(s) matching the latent state, each
          of shape
          `concat([[num_timesteps, num_particles, b1, ..., bN], event_shape])`,
          representing unbiased samples from the posterior distribution
          `p(latent_states | observations)`.
        incremental_log_marginal_likelihoods: float `Tensor` of shape
          `[num_observation_steps, b1, ..., bN]`,
          giving the natural logarithm of an unbiased estimate of
          `p(observations[t] | observations[:t])` at each timestep `t`. Note that
          (by [Jensen's inequality](
          https://en.wikipedia.org/wiki/Jensen%27s_inequality))
          this is *smaller* in expectation than the true
          `log p(observations[t] | observations[:t])`.
  
      #### Examples
  
      **Tracking unknown position and velocity**: Let's consider tracking an object
      moving in a one-dimensional space. We'll define a dynamical system
      by specifying an `initial_state_prior`, a `transition_fn`,
      and `observation_fn`.
  
      The structure of the latent state space is determined by the prior
      distribution. Here, we'll define a state space that includes the object's
      current position and velocity:
  
      ```python
      initial_state_prior = tfd.JointDistributionNamed({
          'position': tfd.Normal(loc=0., scale=1.),
          'velocity': tfd.Normal(loc=0., scale=0.1)})
      ```
  
      The `transition_fn` specifies the evolution of the system. It should
      return a distribution over latent states of the same structure as the prior.
      Here, we'll assume that the position evolves according to the velocity,
      with a small random drift, and the velocity also changes slowly, following
      a random drift:
  
      ```python
      def transition_fn(_, previous_state):
        return tfd.JointDistributionNamed({
            'position': tfd.Normal(
                loc=previous_state['position'] + previous_state['velocity'],
                scale=0.1),
            'velocity': tfd.Normal(loc=previous_state['velocity'], scale=0.01)})
      ```
  
      The `observation_fn` specifies the process by which the system is observed
      at each time step. Let's suppose we observe only a noisy version of the =
      current position.
  
      ```python
        def observation_fn(_, state):
          return tfd.Normal(loc=state['position'], scale=0.1)
      ```
  
      Now let's track our object. Suppose we've been given observations
      corresponding to an initial position of `0.4` and constant velocity of `0.01`:
  
      ```python
      # Generate simulated observations.
      observed_positions = tfd.Normal(loc=tf.linspace(0.4, 0.8, 0.01),
                                      scale=0.1).sample()
  
      # Run particle filtering to sample plausible trajectories.
      (trajectories,  # {'position': [40, 1000], 'velocity': [40, 1000]}
       lps) = tfp.experimental.mcmc.infer_trajectories(
                observations=observed_positions,
                initial_state_prior=initial_state_prior,
                transition_fn=transition_fn,
                observation_fn=observation_fn,
                num_particles=1000)
      ```
  
      For all `i`, `trajectories['position'][:, i]` is a sample from the
      posterior over position sequences, given the observations:
      `p(state[0:T] | observations[0:T])`. Often, the sampled trajectories
      will be highly redundant in their earlier timesteps, because most
      of the initial particles have been discarded through resampling
      (this problem is known as 'particle degeneracy'; see section 3.5 of
      [Doucet and Johansen][1]).
      In such cases it may be useful to also consider the series of *filtering*
      distributions `p(state[t] | observations[:t])`, in which each latent state
      is inferred conditioned only on observations up to that point in time; these
      may be computed using `tfp.mcmc.experimental.particle_filter`.
  
      #### References
  
      [1] Arnaud Doucet and Adam M. Johansen. A tutorial on particle
          filtering and smoothing: Fifteen years later.
          _Handbook of nonlinear filtering_, 12(656-704), 2009.
          https://www.stats.ox.ac.uk/~doucet/doucet_johansen_tutorialPF2011.pdf
      [2] 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 "infer_trajectories") as name:
          pf_seed, resample_seed = samplers.split_seed(
              seed, salt="infer_trajectories"
          )
          (
              particles,
              log_weights,
              parent_indices,
              incremental_log_marginal_likelihoods,
          ) = particle_filter(
              observations=observations,
              initial_state_prior=initial_state_prior,
              transition_fn=transition_fn,
              observation_fn=observation_fn,
              num_particles=num_particles,
              initial_state_proposal=initial_state_proposal,
              particles_dim=particles_dim,
              proposal_fn=proposal_fn,
              resample_fn=resample_fn,
              resample_criterion_fn=resample_criterion_fn,
              unbiased_gradients=unbiased_gradients,
              num_transitions_per_observation=num_transitions_per_observation,
              trace_fn=_default_trace_fn,
              trace_criterion_fn=lambda *_: True,
              seed=pf_seed,
              name=name,
          )
          weighted_trajectories = reconstruct_trajectories(
              particles, parent_indices, particles_dim=particles_dim
          )
  
          # Resample all steps of the trajectories using the final weights.
          resample_indices = resample_fn(
              log_probs=log_weights[-1],
              event_size=num_particles,
              particles_dim=particles_dim,
              sample_shape=(),
              seed=resample_seed,
          )
          trajectories = tf.nest.map_structure(
              lambda x: mcmc_util.index_remapping_gather(
                  x,  # pylint: disable=g-long-lambda
                  resample_indices,
                  axis=particles_dim + 1,
                  indices_axis=particles_dim,
              ),
              weighted_trajectories,
          )
  
          return trajectories, incremental_log_marginal_likelihoods
  
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 2 for calling a global function with example usage: infer_trajectories(observations, initial_state_prior, transition_fn, observation_fn, num_particles, initial_state_proposal, proposal_fn, resample_fn, resample_criterion_fn, unbiased_gradients, num_transitions_per_observation, particles_dim, seed, name) and returns: trajectories,  incremental_log_marginal_likelihoods