code: def test_shapes(
      batch_initial, batch_transition, batch_observation, batch_data
  ):
      # Dimensions.
      batch_size = 3
      state_dim, observation_dim, num_steps = 4, 5, 6
  
      # Model initialization.
      initial_logits = torch.randn(
          [batch_size] * batch_initial + [state_dim]
      )
      initial_logits = initial_logits - initial_logits.logsumexp(
          -1, True
      )
      transition_logits = torch.randn(
          [batch_size] * batch_transition + [state_dim, state_dim]
      )
      transition_logits = (
          transition_logits - transition_logits.logsumexp(-1, True)
      )
      observation_logits = torch.randn(
          [batch_size] * batch_observation
          + [state_dim, observation_dim]
      )
      observation_logits = (
          observation_logits - observation_logits.logsumexp(-1, True)
      )
  
      hmm = MissingDataDiscreteHMM(
          initial_logits, transition_logits, observation_logits
      )
  
      # Random observations.
      value = (
          torch.randint(
              observation_dim, [batch_size] * batch_data + [num_steps]
          ).unsqueeze(-1)
          == torch.arange(observation_dim)
      ).double()
  
      # Log probability.
      lp = hmm.log_prob(value)
  
      # Check shapes:
      if all(
          [
              not batch_initial,
              not batch_transition,
              not batch_observation,
              not batch_data,
          ]
      ):
          assert lp.shape == ()
      else:
          assert lp.shape == (batch_size,)
  
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 2 for calling a global function with example usage: test_shapes(batch_initial, batch_transition, batch_observation, batch_data) without return types