code: def test_log_prob_matches_linear_gaussian_ssm(self):
      dim = 2
      batch_shape = [3, 1]
      seed, *model_seeds = samplers.split_seed(
          test_util.test_seed(), n=6
      )
  
      # Sample a random linear Gaussian process.
      prior_loc = self.evaluate(
          normal.Normal(0.0, 1.0).sample(
              batch_shape + [dim], seed=model_seeds[0]
          )
      )
      prior_scale = self.evaluate(
          inverse_gamma.InverseGamma(1.0, 1.0).sample(
              batch_shape + [dim], seed=model_seeds[1]
          )
      )
      transition_matrix = self.evaluate(
          normal.Normal(0.0, 1.0).sample(
              [dim, dim], seed=model_seeds[2]
          )
      )
      transition_bias = self.evaluate(
          normal.Normal(0.0, 1.0).sample(
              batch_shape + [dim], seed=model_seeds[3]
          )
      )
      transition_scale_tril = self.evaluate(
          tf.linalg.cholesky(
              wishart.WishartTriL(
                  df=dim, scale_tril=tf.eye(dim)
              ).sample(seed=model_seeds[4])
          )
      )
  
      initial_state_prior = mvn_diag.MultivariateNormalDiag(
          loc=prior_loc,
          scale_diag=prior_scale,
          name="initial_state_prior",
      )
  
      lgssm = lgssm_lib.LinearGaussianStateSpaceModel(
          num_timesteps=7,
          transition_matrix=transition_matrix,
          transition_noise=mvn_tril.MultivariateNormalTriL(
              loc=transition_bias, scale_tril=transition_scale_tril
          ),
          # Trivial observation model to pass through the latent state.
          observation_matrix=tf.eye(dim),
          observation_noise=mvn_diag.MultivariateNormalDiag(
              loc=tf.zeros(dim), scale_diag=tf.zeros(dim)
          ),
          initial_state_prior=initial_state_prior,
      )
  
      markov_chain = markov_chain_lib.MarkovChain(
          initial_state_prior=initial_state_prior,
          transition_fn=lambda _, x: mvn_tril.MultivariateNormalTriL(  # pylint: disable=g-long-lambda
              loc=tf.linalg.matvec(transition_matrix, x)
              + transition_bias,
              scale_tril=transition_scale_tril,
          ),
          num_steps=7,
      )
  
      x = markov_chain.sample(5, seed=seed)
      self.assertAllClose(
          lgssm.log_prob(x), markov_chain.log_prob(x), rtol=1e-5
      )
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import split
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import inverse_gamma
  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 log_prob_ratio
  from tensorflow_probability.python.distributions import markov_chain
  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 uniform
  from tensorflow_probability.python.distributions import wishart
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 2 for calling the MarkovChainTest obj's underlying member method code with example usage: obj.test_log_prob_matches_linear_gaussian_ssm() without return types