code: def test_sample_given_states():
      a0 = torch.tensor([0.9, 0.08, 0.02])
      a = torch.tensor(
          [[0.1, 0.8, 0.1], [0.5, 0.3, 0.2], [0.4, 0.4, 0.2]]
      )
      eps = 1e-10
      # Effectively deterministic to check sampler.
      e = torch.tensor([[1 - eps, eps], [eps, 1 - eps], [eps, 1 - eps]])
  
      map_states = torch.tensor([0, 2, 1, 0], dtype=torch.long)
      hmm_distr = MissingDataDiscreteHMM(
          torch.log(a0), torch.log(a), torch.log(e)
      )
      sample = hmm_distr.sample_given_states(map_states)
      expected_sample = torch.tensor(
          [[1.0, 0.0], [0.0, 1.0], [0.0, 1.0], [1.0, 0.0]]
      )
      assert torch.allclose(sample, expected_sample)
  
      # Batch values
      map_states = torch.tensor(
          [[0, 2, 1, 0], [0, 0, 0, 1]], dtype=torch.long
      )
      sample = hmm_distr.sample_given_states(map_states)
      expected_sample = torch.tensor(
          [
              [[1.0, 0.0], [0.0, 1.0], [0.0, 1.0], [1.0, 0.0]],
              [[1.0, 0.0], [1.0, 0.0], [1.0, 0.0], [0.0, 1.0]],
          ]
      )
      assert torch.allclose(sample, expected_sample)
  
      # Batch parameters
      e = torch.cat(
          [
              e[None, :, :],
              torch.tensor(
                  [[eps, 1 - eps], [eps, 1 - eps], [1 - eps, eps]]
              )[None, :, :],
          ],
          dim=0,
      )
      hmm_distr = MissingDataDiscreteHMM(
          torch.log(a0), torch.log(a), torch.log(e)
      )
      sample = hmm_distr.sample_given_states(map_states[0])
      expected_sample = torch.tensor(
          [
              [[1.0, 0.0], [0.0, 1.0], [0.0, 1.0], [1.0, 0.0]],
              [[0.0, 1.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]],
          ]
      )
      assert torch.allclose(sample, expected_sample)
  
      # Batch parameters and values.
      sample = hmm_distr.sample_given_states(map_states)
      expected_sample = torch.tensor(
          [
              [[1.0, 0.0], [0.0, 1.0], [0.0, 1.0], [1.0, 0.0]],
              [[0.0, 1.0], [0.0, 1.0], [0.0, 1.0], [0.0, 1.0]],
          ]
      )
      assert torch.allclose(sample, expected_sample)
  
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 9 for calling a global function with example usage: test_sample_given_states() without return types