code: def test_switching_linear_hmm_log_prob(
      exact, num_steps, hidden_dim, obs_dim, num_components
  ):
      # This tests agreement between an SLDS and an HMM when all components
      # are identical, i.e. so latent can be marginalized out.
      torch.manual_seed(2)
      init_logits = torch.rand(num_components)
      init_mvn = random_mvn((), hidden_dim)
      trans_logits = torch.rand(num_components)
      trans_matrix = torch.randn(hidden_dim, hidden_dim)
      trans_mvn = random_mvn((), hidden_dim)
      obs_matrix = torch.randn(hidden_dim, obs_dim)
      obs_mvn = random_mvn((), obs_dim)
  
      expected_dist = GaussianHMM(
          init_mvn,
          trans_matrix.expand(num_steps, -1, -1),
          trans_mvn,
          obs_matrix,
          obs_mvn,
      )
      actual_dist = SwitchingLinearHMM(
          init_logits,
          init_mvn,
          trans_logits,
          trans_matrix.expand(num_steps, num_components, -1, -1),
          trans_mvn,
          obs_matrix,
          obs_mvn,
          exact=exact,
      )
      assert actual_dist.batch_shape == expected_dist.batch_shape
      assert actual_dist.event_shape == expected_dist.event_shape
  
      data = obs_mvn.sample(expected_dist.batch_shape + (num_steps,))
      assert data.shape == expected_dist.shape()
      expected_log_prob = expected_dist.log_prob(data)
      assert expected_log_prob.shape == expected_dist.batch_shape
      actual_log_prob = actual_dist.log_prob(data)
      assert_close(
          actual_log_prob, expected_log_prob, atol=1e-4, rtol=None
      )
  
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 9 for calling a global function with example usage: test_switching_linear_hmm_log_prob(exact, num_steps, hidden_dim, obs_dim, num_components) without return types