code: def test_custom_trace_fn(self):
      def trace_fn(state, _):
          # Traces the mean and stddev of the particle population at each step.
          weights = tf.exp(state.log_weights)
          mean = tf.reduce_sum(weights * state.particles, axis=0)
          variance = tf.reduce_sum(
              weights * (state.particles - mean[tf.newaxis, ...]) ** 2
          )
          return {
              "mean": mean,
              "stddev": tf.sqrt(variance),
              # In real usage we would likely not track the particles and
              # weights. We keep them here just so we can double-check the
              # stats, below.
              "particles": state.particles,
              "weights": weights,
          }
  
      results = 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, 1.0),
              observation_fn=lambda _, state: normal.Normal(state, 1.0),
              num_particles=1024,
              trace_fn=trace_fn,
              seed=test_util.test_seed(),
          )
      )
  
      # Verify that posterior means are increasing.
      self.assertAllGreater(
          results["mean"][1:] - results["mean"][:-1], 0.0
      )
  
      # Check that our traced means and scales match values computed
      # by averaging over particles after the fact.
      all_means = self.evaluate(
          tf.reduce_sum(
              results["weights"] * results["particles"], axis=1
          )
      )
      all_variances = self.evaluate(
          tf.reduce_sum(
              results["weights"]
              * (results["particles"] - all_means[..., tf.newaxis])
              ** 2,
              axis=1,
          )
      )
      self.assertAllClose(results["mean"], all_means)
      self.assertAllClose(results["stddev"], np.sqrt(all_variances))
  
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 10 for calling the _ParticleFilterTest obj's underlying member method code with example usage: obj.test_custom_trace_fn() and returns: dict