code: def test_particle_tree(self):
      # m particles are born at the same random location on an n x n grid.
      # They independently take `n_steps` steps of a random walk going N, S,
      # E or W at each step. At the end we observe their positions subject
      # to Gaussian noise. Computes posterior distribution for the birthplace.
  
      n = 16
      m = 3
      n_steps = 16
  
      def model():
          # Shared birthplace
          x_start = yield Root(
              categorical.Categorical(
                  probs=tf.ones(n) / n, dtype=tf.int32
              )
          )
          y_start = yield categorical.Categorical(
              probs=tf.ones(n) / n, dtype=tf.int32
          )
  
          x = m * [x_start]
          y = m * [y_start]
  
          for t in range(n_steps):
              for i in range(m):
                  # Construct PDF for next step in walk
                  # Start with PDF for all mass on current point.
                  ox = tf.one_hot(x[i], n)
                  oy = tf.one_hot(y[i], n)
                  o = ox[..., :, None] * oy[..., None, :]
  
                  # Deliberate choice of non-centered distribution as
                  # reduced symmetry lends itself to better testing.
                  p = (
                      0.1 * tf.roll(o, shift=[0, -1], axis=[-2, -1])
                      + 0.2 * tf.roll(o, shift=[0, 1], axis=[-2, -1])
                      + 0.3 * tf.roll(o, shift=[-1, 0], axis=[-2, -1])
                      + 0.4 * tf.roll(o, shift=[1, 0], axis=[-2, -1])
                  )
  
                  # Reshape just last two dimensions.
                  p = tf.reshape(
                      p, ps.concat([ps.shape(p)[:-2], [-1]], axis=0)
                  )
                  xy = yield categorical.Categorical(
                      probs=p, dtype=tf.int32
                  )
                  x[i] = xy // n
                  y[i] = xy % n
  
          # 2 * m noisy 2D observations at end
          for i in range(m):
              yield normal.Normal(
                  tf.cast(x[i], dtype=tf.float32), scale=2.0
              )
              yield normal.Normal(
                  tf.cast(y[i], dtype=tf.float32), scale=2.0
              )
  
      d = marginalize.MarginalizableJointDistributionCoroutine(model)
      final_observations = [0.0, 0.0, 4.0, 4.0, 0.0, 4.0]
      observations = (
          ["tabulate", "tabulate"]
          + n_steps * ["marginalize", "marginalize", "marginalize"]
          + final_observations
      )
      # Using `method='logeinsumexp' here creates a large intermediate
      # that impacts performance. Because we are only setting
      # `STEPS == 16` the probabilities involved aren't small enough
      # to result in underflow.
      p = tf.exp(d.marginalized_log_prob(observations, method="einsum"))
      q = _tree_example(n, n_steps)
  
      # Note that while p and q should be close in value there is a large
      # difference in computation time. I would expect the p
      # to be slower by a factor of around `n_steps/log2(n_steps)` (because
      # `numpy` compute matrix powers by repeated squaring) but it seems
      # to be even slower. This likely means future versions of
      # `marginalized_log_prob` will run faster when the elimination
      # order chosen by `tf.einsum` closer matches `_tree_example` above.
      self.assertAllClose(p, q)
  
imports: import functools
  import itertools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import hidden_markov_model
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  import tensorflow_probability.python.experimental.marginalize
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 7 for calling the _MarginalizeTest obj's underlying member method code with example usage: obj.test_particle_tree() without return types