code: def test_switching_linear_hmm_shape(
      init_cat_shape,
      init_mvn_shape,
      trans_cat_shape,
      trans_mat_shape,
      trans_mvn_shape,
      obs_mat_shape,
      obs_mvn_shape,
  ):
      hidden_dim, obs_dim = obs_mat_shape[-2:]
      assert trans_mat_shape[-2:] == (hidden_dim, hidden_dim)
  
      init_logits = torch.randn(init_cat_shape)
      init_mvn = random_mvn(init_mvn_shape, hidden_dim)
      trans_logits = torch.randn(trans_cat_shape)
      trans_matrix = torch.randn(trans_mat_shape)
      trans_mvn = random_mvn(trans_mvn_shape, hidden_dim)
      obs_matrix = torch.randn(obs_mat_shape)
      obs_mvn = random_mvn(obs_mvn_shape, obs_dim)
  
      init_shape = broadcast_shape(init_cat_shape, init_mvn_shape)
      shape = broadcast_shape(
          init_shape[:-1] + (1, init_shape[-1]),
          trans_cat_shape[:-1],
          trans_mat_shape[:-2],
          trans_mvn_shape,
          obs_mat_shape[:-2],
          obs_mvn_shape,
      )
      expected_batch_shape, time_shape = shape[:-2], shape[-2:-1]
      expected_event_shape = time_shape + (obs_dim,)
  
      actual_dist = SwitchingLinearHMM(
          init_logits,
          init_mvn,
          trans_logits,
          trans_matrix,
          trans_mvn,
          obs_matrix,
          obs_mvn,
      )
      assert actual_dist.event_shape == expected_event_shape
      assert actual_dist.batch_shape == expected_batch_shape
  
      data = obs_mvn.expand(shape).sample()[..., 0, :]
      actual_log_prob = actual_dist.log_prob(data)
      assert actual_log_prob.shape == expected_batch_shape
      check_expand(actual_dist, data)
  
      final_cat, final_mvn = actual_dist.filter(data)
      assert isinstance(final_cat, dist.Categorical)
      assert isinstance(final_mvn, dist.MultivariateNormal)
      assert final_cat.batch_shape == actual_dist.batch_shape
      assert (
          final_mvn.batch_shape
          == actual_dist.batch_shape + final_cat.logits.shape[-1:]
      )
  
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 8 for calling a global function with example usage: test_switching_linear_hmm_shape(init_cat_shape, init_mvn_shape, trans_cat_shape, trans_mat_shape, trans_mvn_shape, obs_mat_shape, obs_mvn_shape) without return types