code: def test_discrete_mvn_log_prob(
      init_shape, trans_shape, obs_shape, state_dim
  ):
      event_size = 4
      init_logits = torch.randn(init_shape + (state_dim,))
      trans_logits = torch.randn(trans_shape + (state_dim, state_dim))
      loc = torch.randn(obs_shape + (state_dim, event_size))
      cov = torch.randn(
          obs_shape + (state_dim, event_size, 2 * event_size)
      )
      cov = cov.matmul(cov.transpose(-1, -2))
      scale_tril = torch.cholesky(cov)
      obs_dist = dist.MultivariateNormal(loc, scale_tril=scale_tril)
  
      actual_dist = DiscreteHMM(init_logits, trans_logits, obs_dist)
      expected_dist = dist.DiscreteHMM(
          init_logits, trans_logits, obs_dist
      )
      assert actual_dist.event_shape == expected_dist.event_shape
      assert actual_dist.batch_shape == expected_dist.batch_shape
  
      batch_shape = broadcast_shape(
          init_shape + (1,), trans_shape, obs_shape
      )
      data = obs_dist.expand(batch_shape + (state_dim,)).sample()
      data = data[(slice(None),) * len(batch_shape) + (0,)]
      actual_log_prob = actual_dist.log_prob(data)
      expected_log_prob = expected_dist.log_prob(data)
      assert_close(actual_log_prob, expected_log_prob)
      check_expand(actual_dist, data)
  
file_path: funsor/test/pyro/test_hmm.py
imports: import pyro.distributions
  import pytest
  import torch
  from funsor.pyro.hmm import DiscreteHMM, GaussianHMM, GaussianMRF, SwitchingLinearHMM
  from funsor.testing import assert_close, random_mvn
  from funsor.util import broadcast_shape
label_desc: Use image node_id 4 for calling a global function with example usage: test_discrete_mvn_log_prob(init_shape, trans_shape, obs_shape, state_dim) without return types