code: def test_discrete_hmm_homogeneous_trick(
      init_shape,
      trans_shape,
      obs_shape,
      event_shape,
      state_dim,
      num_steps,
  ):
      batch_shape = broadcast_shape(
          init_shape, trans_shape[:-1], obs_shape[:-1]
      )
      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)
      )
  
      d = dist.DiscreteHMM(init_logits, trans_logits, obs_dist)
      assert d.event_shape == (1,) + event_shape
      assert d.support.event_dim == d.event_dim
  
      data = obs_dist.expand(
          batch_shape + (num_steps, state_dim)
      ).sample()
      data = data[(slice(None),) * (len(batch_shape) + 1) + (0,)]
      assert data.shape == batch_shape + (num_steps,) + event_shape
      actual = d.log_prob(data)
      assert actual.shape == batch_shape
  
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 6 for calling a global function with example usage: test_discrete_hmm_homogeneous_trick(init_shape, trans_shape, obs_shape, event_shape, state_dim, num_steps) without return types