code: def test_discrete_hmm_shape(
      ok, init_shape, trans_shape, obs_shape, event_shape, state_dim
  ):
      init_logits = torch.randn(init_shape + (state_dim,))
      trans_logits = torch.randn(trans_shape + (state_dim, state_dim))
      obs_logits = torch.randn(obs_shape + (state_dim,) + event_shape)
      obs_dist = dist.Bernoulli(logits=obs_logits).to_event(
          len(event_shape)
      )
      data = obs_dist.sample()[(slice(None),) * len(obs_shape) + (0,)]
  
      if not ok:
          with pytest.raises(ValueError):
              d = dist.DiscreteHMM(init_logits, trans_logits, obs_dist)
              d.log_prob(data)
          return
  
      d = dist.DiscreteHMM(init_logits, trans_logits, obs_dist)
      assert d.support.event_dim == d.event_dim
  
      actual = d.log_prob(data)
      expected_shape = broadcast_shape(
          init_shape, trans_shape[:-1], obs_shape[:-1]
      )
      assert actual.shape == expected_shape
      check_expand(d, data)
      check_sample_shape(d)
  
      final = d.filter(data)
      assert isinstance(final, dist.Categorical)
      assert final.batch_shape == d.batch_shape
      assert final.event_shape == ()
      assert final.support.upper_bound == state_dim - 1
  
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 5 for calling a global function with example usage: test_discrete_hmm_shape(ok, init_shape, trans_shape, obs_shape, event_shape, state_dim) without return types