code: def _tree_example(n, n_steps):
      # Computes result for `test_particle_tree` using
      # explicit "hand-constructed" computation based on
      # matrix multiplication.
  
      init = np.ones([n * n]) / (n * n)
      trans = np.zeros((n * n, n * n), dtype=np.float64)
  
      def wrap(i):
          return (i + n) % n
  
      def index(i, j):
          return wrap(i) * n + wrap(j)
  
      # Compute transition matrix for random walk.
      for i in range(n):
          for j in range(n):
              trans[index(i, j), index(i, j - 1)] = 0.1
              trans[index(i, j), index(i, j + 1)] = 0.2
              trans[index(i, j), index(i - 1, j)] = 0.3
              trans[index(i, j), index(i + 1, j)] = 0.4
  
      def normal_prob(loc, scale, z):
          return (
              np.exp(-((z - loc) ** 2) / (2.0 * scale * scale))
              / np.sqrt(2 * np.pi)
              / scale
          )
  
      # `n_steps` steps in random walk.
      trans10 = np.linalg.matrix_power(trans, n_steps)
  
      zs = np.arange(0, n)
  
      def obs(x, y):
          return (
              normal_prob(zs, 2.0, x)[:, None] * normal_prob(zs, 2.0, y)
          ).reshape([n * n])
  
      v1 = obs(0.0, 0.0)
      v2 = obs(4.0, 4.0)
      v3 = obs(0.0, 4.0)
  
      p1 = np.matmul(trans10, v1)
      p2 = np.matmul(trans10, v2)
      p3 = np.matmul(trans10, v3)
  
      p = (init * p1 * p2 * p3).reshape([n, n])
  
      return p
  
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 3 for calling a global function with example usage: _tree_example(n, n_steps) and returns: p, unknown, unknown, unknown, unknown