code: def test_discrete_hmm_categorical(num_steps):
      state_dim = 3
      obs_dim = 4
      init_logits = torch.randn(state_dim)
      trans_logits = torch.randn(num_steps, state_dim, state_dim)
      obs_dist = dist.Categorical(
          logits=torch.randn(num_steps, state_dim, obs_dim)
      )
      d = dist.DiscreteHMM(init_logits, trans_logits, obs_dist)
      data = dist.Categorical(
          logits=torch.zeros(num_steps, obs_dim)
      ).sample()
      actual = d.log_prob(data)
      assert actual.shape == d.batch_shape
      check_expand(d, data)
      check_sample_shape(d)
  
      # Check loss against TraceEnum_ELBO.
      @config_enumerate
      def model(data):
          x = pyro.sample(
              "x_init", dist.Categorical(logits=init_logits)
          )
          for t in range(num_steps):
              x = pyro.sample(
                  "x_{}".format(t),
                  dist.Categorical(
                      logits=Vindex(trans_logits)[..., t, x, :]
                  ),
              )
              pyro.sample(
                  "obs_{}".format(t),
                  dist.Categorical(
                      logits=Vindex(obs_dist.logits)[..., t, x, :]
                  ),
                  obs=data[..., t],
              )
  
      expected_loss = TraceEnum_ELBO().loss(model, empty_guide, data)
      actual_loss = -float(actual.sum())
      assert_close(actual_loss, expected_loss)
  
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 8 for calling a global function with example usage: test_discrete_hmm_categorical(num_steps) without return types