code: def test_sample_states():
      # Effectively deterministic to check sampler.
      eps = 1e-10
      a0 = torch.tensor([1 - eps, eps / 2, eps / 2])
      a = torch.tensor(
          [
              [eps / 2, 1 - eps, eps / 2],
              [eps, 0.5 - eps / 2, 0.5 - eps / 2],
              [eps, 0.5 - eps / 2, 0.5 - eps / 2],
          ]
      )
      e = torch.tensor([[1 - eps, eps], [1 - eps, eps], [eps, 1 - eps]])
      x = torch.tensor([[1.0, 0.0], [0.0, 0.0], [0.0, 1.0], [0.0, 1.0]])
  
      hmm_distr = MissingDataDiscreteHMM(
          torch.log(a0), torch.log(a), torch.log(e)
      )
      states = hmm_distr.sample_states(x)
      expected_states = torch.tensor([0, 1, 2, 2])
      assert torch.allclose(states, expected_states)
  
      # Batch values.
      x = torch.cat(
          [
              x[None, :, :],
              torch.tensor(
                  [[1.0, 0.0], [0.0, 0.0], [0.0, 1.0], [1.0, 0.0]]
              )[None, :, :],
          ],
          dim=0,
      )
      states = hmm_distr.sample_states(x)
      expected_states = torch.tensor([[0, 1, 2, 2], [0, 1, 2, 1]])
      assert torch.allclose(states, expected_states)
  
      # Batch parameters
      a0 = torch.cat(
          [
              a0[None, :],
              torch.tensor([eps / 2, 1 - eps, eps / 2])[None, :],
          ]
      )
      a = torch.cat(
          [
              a[None, :, :],
              torch.tensor(
                  [
                      [eps / 2, 1 - eps, eps / 2],
                      [eps / 2, 1 - eps, eps / 2],
                      [eps / 2, 1 - eps, eps / 2],
                  ]
              )[None, :, :],
          ],
          dim=0,
      )
      e = torch.cat(
          [
              e[None, :, :],
              torch.tensor(
                  [[1 - eps, eps], [0.5, 0.5], [eps, 1 - eps]]
              )[None, :, :],
          ],
          dim=0,
      )
      hmm_distr = MissingDataDiscreteHMM(
          torch.log(a0), torch.log(a), torch.log(e)
      )
      states = hmm_distr.sample_states(x[1])
      expected_states = torch.tensor([[0, 1, 2, 1], [1, 1, 1, 1]])
      assert torch.allclose(states, expected_states)
  
      # Batch both parameters and values.
      states = hmm_distr.sample_states(x)
      expected_states = torch.tensor([[0, 1, 2, 2], [1, 1, 1, 1]])
      assert torch.allclose(states, expected_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 10 for calling a global function with example usage: test_sample_states() without return types