code: def test_gaussian_hmm_distribution(
      diag, sample_shape, batch_shape, num_steps, hidden_dim, obs_dim
  ):
      init_dist = random_mvn(batch_shape, hidden_dim)
      trans_mat = torch.randn(
          batch_shape + (num_steps, hidden_dim, hidden_dim)
      )
      trans_dist = random_mvn(batch_shape + (num_steps,), hidden_dim)
      obs_mat = torch.randn(
          batch_shape + (num_steps, hidden_dim, obs_dim)
      )
      obs_dist = random_mvn(batch_shape + (num_steps,), obs_dim)
      if diag:
          scale = obs_dist.scale_tril.diagonal(dim1=-2, dim2=-1)
          obs_dist = dist.Normal(obs_dist.loc, scale).to_event(1)
      d = dist.GaussianHMM(
          init_dist,
          trans_mat,
          trans_dist,
          obs_mat,
          obs_dist,
          duration=num_steps,
      )
      if diag:
          obs_mvn = dist.MultivariateNormal(
              obs_dist.base_dist.loc,
              scale_tril=obs_dist.base_dist.scale.diag_embed(),
          )
      else:
          obs_mvn = obs_dist
      data = obs_dist.sample(sample_shape)
      assert data.shape == sample_shape + d.shape()
      actual_log_prob = d.log_prob(data)
  
      # Compare against hand-computed density.
      # We will construct enormous unrolled joint gaussians with shapes:
      #       t | 0 1 2 3 1 2 3      T = 3 in this example
      #   ------+-----------------------------------------
      #    init | H
      #   trans | H H H H            H = hidden
      #     obs |   H H H O O O      O = observed
      #    like |         O O O
      # and then combine these using gaussian_tensordot().
      T = num_steps
      init = mvn_to_gaussian(init_dist)
      trans = matrix_and_mvn_to_gaussian(trans_mat, trans_dist)
      obs = matrix_and_mvn_to_gaussian(obs_mat, obs_mvn)
      like_dist = dist.Normal(torch.randn(data.shape), 1).to_event(2)
      like = mvn_to_gaussian(like_dist)
  
      unrolled_trans = reduce(
          operator.add,
          [
              trans[..., t].event_pad(
                  left=t * hidden_dim, right=(T - t - 1) * hidden_dim
              )
              for t in range(T)
          ],
      )
      unrolled_obs = reduce(
          operator.add,
          [
              obs[..., t].event_pad(
                  left=t * obs.dim(), right=(T - t - 1) * obs.dim()
              )
              for t in range(T)
          ],
      )
      unrolled_like = reduce(
          operator.add,
          [
              like[..., t].event_pad(
                  left=t * obs_dim, right=(T - t - 1) * obs_dim
              )
              for t in range(T)
          ],
      )
      # Permute obs from HOHOHO to HHHOOO.
      perm = torch.cat(
          [torch.arange(hidden_dim) + t * obs.dim() for t in range(T)]
          + [
              torch.arange(obs_dim) + hidden_dim + t * obs.dim()
              for t in range(T)
          ]
      )
      unrolled_obs = unrolled_obs.event_permute(perm)
      unrolled_data = data.reshape(data.shape[:-2] + (T * obs_dim,))
  
      assert init.dim() == hidden_dim
      assert unrolled_trans.dim() == (1 + T) * hidden_dim
      assert unrolled_obs.dim() == T * (hidden_dim + obs_dim)
      logp = gaussian_tensordot(init, unrolled_trans, hidden_dim)
      logp = gaussian_tensordot(logp, unrolled_obs, T * hidden_dim)
      expected_log_prob = logp.log_density(unrolled_data)
      assert_close(actual_log_prob, expected_log_prob)
  
      d_posterior, log_normalizer = d.conjugate_update(like_dist)
      assert_close(
          d.log_prob(data) + like_dist.log_prob(data),
          d_posterior.log_prob(data) + log_normalizer,
      )
  
      if batch_shape or sample_shape:
          return
  
      # Test mean and covariance.
      prior = "prior", d, logp
      posterior = "posterior", d_posterior, logp + unrolled_like
      for name, d, g in [prior, posterior]:
          logging.info("testing {} moments".format(name))
          with torch.no_grad():
              num_samples = 100000
              samples = d.sample([num_samples]).reshape(
                  num_samples, T * obs_dim
              )
              actual_mean = samples.mean(0)
              delta = samples - actual_mean
              actual_cov = (
                  delta.unsqueeze(-1) * delta.unsqueeze(-2)
              ).mean(0)
              actual_std = actual_cov.diagonal(dim1=-2, dim2=-1).sqrt()
              actual_corr = actual_cov / (
                  actual_std.unsqueeze(-1) * actual_std.unsqueeze(-2)
              )
  
              expected_cov = torch.linalg.cholesky(
                  g.precision
              ).cholesky_inverse()
              expected_mean = expected_cov.matmul(
                  g.info_vec.unsqueeze(-1)
              ).squeeze(-1)
              expected_std = expected_cov.diagonal(
                  dim1=-2, dim2=-1
              ).sqrt()
              expected_corr = expected_cov / (
                  expected_std.unsqueeze(-1)
                  * expected_std.unsqueeze(-2)
              )
  
              assert_close(
                  actual_mean, expected_mean, atol=0.05, rtol=0.02
              )
              assert_close(
                  actual_std, expected_std, atol=0.05, rtol=0.02
              )
              assert_close(actual_corr, expected_corr, atol=0.02)
  
imports: import logging
  import operator
  from functools import reduce
  import opt_einsum
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.distributions.hmm import _sequential_gamma_gaussian_tensordot, _sequential_logmatmulexp
  from pyro.distributions.util import broadcast_shape
  from pyro.infer import TraceEnum_ELBO, config_enumerate
  from pyro.ops.gamma_gaussian import gamma_and_mvn_to_gamma_gaussian, gamma_gaussian_tensordot, matrix_and_mvn_to_gamma_gaussian
  from pyro.ops.gaussian import gaussian_tensordot, matrix_and_mvn_to_gaussian, mvn_to_gaussian
  from pyro.ops.indexing import Vindex
  from tests.common import assert_close
  from tests.ops.gamma_gaussian import assert_close_gamma_gaussian, random_gamma, random_gamma_gaussian
  from tests.ops.gaussian import random_mvn
label_desc: Use image node_id 13 for calling a global function with example usage: test_gaussian_hmm_distribution(diag, sample_shape, batch_shape, num_steps, hidden_dim, obs_dim) without return types