code: def test_hmm(self):
      n_steps = 4
      infer_step = 2
  
      observations = np.array([-1.0, 0.0, 1.0, 2.0], np.float32)
  
      initial_prob = tf.constant([0.6, 0.4], dtype=tf.float32)
      transition_matrix = tf.constant(
          [[0.6, 0.4], [0.3, 0.7]], dtype=tf.float32
      )
      observation_locs = tf.constant([0.0, 1.0], dtype=tf.float32)
      observation_scale = tf.constant(0.5, dtype=tf.float32)
  
      dist1 = hidden_markov_model.HiddenMarkovModel(
          categorical.Categorical(probs=initial_prob),
          categorical.Categorical(probs=transition_matrix),
          normal.Normal(loc=observation_locs, scale=observation_scale),
          num_steps=n_steps,
      )
  
      p = dist1.posterior_marginals(observations).probs_parameter()[
          infer_step
      ]
  
      def model():
          i = yield Root(
              categorical.Categorical(
                  probs=initial_prob, dtype=tf.int32
              )
          )
          z = yield normal.Normal(
              loc=tf.gather(observation_locs, i),
              scale=observation_scale,
          )
  
          for t in range(n_steps - 1):
              i = yield categorical.Categorical(
                  probs=tf.gather(transition_matrix, i), dtype=tf.int32
              )
              yield normal.Normal(
                  loc=tf.gather(observation_locs, i),
                  scale=observation_scale,
              )
  
      dist2 = marginalize.MarginalizableJointDistributionCoroutine(
          model
      )
      full_observations = list(
          itertools.chain(
              *zip(
                  [
                      "tabulate" if i == infer_step else "marginalize"
                      for i in range(n_steps)
                  ],
                  observations,
              )
          )
      )
      q = tf.exp(dist2.marginalized_log_prob(full_observations))
      q = q / tf.reduce_sum(q)
  
      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 5 for calling the _MarginalizeTest obj's underlying member method code with example usage: obj.test_hmm() without return types