code: def test_step_indices_to_trace(self):
      num_particles = 1024
      (
          particles_1_3,
          log_weights_1_3,
          parent_indices_1_3,
          incremental_log_marginal_likelihood_1_3,
      ) = self.evaluate(
          particle_filter.particle_filter(
              observations=tf.convert_to_tensor(
                  [1.0, 3.0, 5.0, 7.0, 9.0]
              ),
              initial_state_prior=normal.Normal(0.0, 1.0),
              transition_fn=lambda _, state: normal.Normal(state, 10.0),
              observation_fn=lambda _, state: normal.Normal(state, 0.1),
              num_particles=num_particles,
              trace_criterion_fn=lambda s, r: ps.logical_or(  # pylint: disable=g-long-lambda
                  ps.equal(r.steps, 2), ps.equal(r.steps, 4)
              ),
              static_trace_allocation_size=2,
              seed=test_util.test_seed(),
          )
      )
      self.assertLen(particles_1_3, 2)
      self.assertLen(log_weights_1_3, 2)
      self.assertLen(parent_indices_1_3, 2)
      self.assertLen(incremental_log_marginal_likelihood_1_3, 2)
      means = np.sum(np.exp(log_weights_1_3) * particles_1_3, axis=1)
      self.assertAllClose(means, [3.0, 7.0], atol=1.0)
  
      (
          final_particles,
          final_log_weights,
          final_cumulative_lp,
      ) = self.evaluate(
          particle_filter.particle_filter(
              observations=tf.convert_to_tensor(
                  [1.0, 3.0, 5.0, 7.0, 9.0]
              ),
              initial_state_prior=normal.Normal(0.0, 1.0),
              transition_fn=lambda _, state: normal.Normal(state, 10.0),
              observation_fn=lambda _, state: normal.Normal(state, 0.1),
              num_particles=num_particles,
              trace_fn=lambda s, r: (  # pylint: disable=g-long-lambda
                  s.particles,
                  s.log_weights,
                  r.accumulated_log_marginal_likelihood,
              ),
              trace_criterion_fn=None,
              seed=test_util.test_seed(),
          )
      )
      self.assertLen(final_particles, num_particles)
      self.assertLen(final_log_weights, num_particles)
      self.assertEqual(final_cumulative_lp.shape, ())
      means = np.sum(np.exp(final_log_weights) * final_particles)
      self.assertAllClose(means, 9.0, atol=1.5)
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.mcmc import particle_filter
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 11 for calling the _ParticleFilterTest obj's underlying member method code with example usage: obj.test_step_indices_to_trace() without return types