code: def test_DiscreteHMM_comparison(
      batch_initial, batch_transition, batch_observation, batch_data
  ):
      # Dimensions.
      batch_size = 3
      state_dim, observation_dim, num_steps = 4, 5, 6
  
      # -- Model setup --.
      transition_logits_vldhmm = torch.randn(
          [batch_size] * batch_transition + [state_dim, state_dim]
      )
      transition_logits_vldhmm = (
          transition_logits_vldhmm
          - transition_logits_vldhmm.logsumexp(-1, True)
      )
      # Adjust for DiscreteHMM broadcasting convention.
      transition_logits_dhmm = transition_logits_vldhmm.unsqueeze(-3)
      # Convert between discrete HMM convention for initial state and variable
      # length HMM convention.
      initial_logits_dhmm = torch.randn(
          [batch_size] * batch_initial + [state_dim]
      )
      initial_logits_dhmm = (
          initial_logits_dhmm - initial_logits_dhmm.logsumexp(-1, True)
      )
      initial_logits_vldhmm = (
          initial_logits_dhmm.unsqueeze(-1) + transition_logits_vldhmm
      ).logsumexp(-2)
      observation_logits = torch.randn(
          [batch_size] * batch_observation
          + [state_dim, observation_dim]
      )
      observation_logits = (
          observation_logits - observation_logits.logsumexp(-1, True)
      )
      # Create distribution object for DiscreteHMM
      observation_dist = Categorical(
          logits=observation_logits.unsqueeze(-3)
      )
  
      vldhmm = MissingDataDiscreteHMM(
          initial_logits_vldhmm,
          transition_logits_vldhmm,
          observation_logits,
      )
      dhmm = DiscreteHMM(
          initial_logits_dhmm, transition_logits_dhmm, observation_dist
      )
  
      # Random observations.
      value = torch.randint(
          observation_dim, [batch_size] * batch_data + [num_steps]
      )
      value_oh = (
          value.unsqueeze(-1) == torch.arange(observation_dim)
      ).double()
  
      # -- Check. --
      # Log probability.
      lp_vldhmm = vldhmm.log_prob(value_oh)
      lp_dhmm = dhmm.log_prob(value)
      # Shapes.
      if all(
          [
              not batch_initial,
              not batch_transition,
              not batch_observation,
              not batch_data,
          ]
      ):
          assert lp_vldhmm.shape == ()
      else:
          assert lp_vldhmm.shape == (batch_size,)
      # Values.
      assert torch.allclose(lp_vldhmm, lp_dhmm)
      # Filter.
      filter_dhmm = dhmm.filter(value)
      filter_vldhmm = vldhmm.filter(value_oh)
      assert torch.allclose(
          filter_dhmm.logits, filter_vldhmm[..., -1, :]
      )
      # Check other computations run.
      vldhmm.sample(value_oh.shape[:-1])
      vldhmm.smooth(value_oh)
      vldhmm.sample_states(value_oh)
      map_states = vldhmm.map_states(value_oh)
      print(value_oh.shape, map_states.shape)
      vldhmm.sample_given_states(map_states)
  
file_path: pyro/tests/contrib/mue/test_missingdatahmm.py
imports: import pytest
  import torch
  from pyro.contrib.mue.missingdatahmm import MissingDataDiscreteHMM
  from pyro.distributions import Categorical, DiscreteHMM
label_desc: Use image node_id 3 for calling a global function with example usage: test_DiscreteHMM_comparison(batch_initial, batch_transition, batch_observation, batch_data) without return types