code: def test_state_infer():
      # HMM parameters.
      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]]
      )
      e = torch.tensor([[0.9, 0.1], [0.1, 0.9], [0.5, 0.5]])
      # Observed value.
      x = torch.tensor(
          [
              [0.0, 1.0],
              [1.0, 0.0],
              [0.0, 1.0],
              [0.0, 1.0],
              [1.0, 0.0],
              [0.0, 0.0],
          ]
      )
  
      expected_map_states = indiv_map_states(a0, a, e, x)
      expected_filter = indiv_filter(a0, a, e, x)
      expected_smooth = indiv_smooth(a0, a, e, x)
  
      hmm_distr = MissingDataDiscreteHMM(
          torch.log(a0), torch.log(a), torch.log(e)
      )
      map_states = hmm_distr.map_states(x)
      filter = hmm_distr.filter(x)
      smooth = hmm_distr.smooth(x)
  
      assert torch.allclose(map_states, expected_map_states)
      assert torch.allclose(filter, expected_filter)
      assert torch.allclose(smooth, expected_smooth)
  
      # Batch values.
      x = torch.cat(
          [
              x[None, :, :],
              torch.tensor(
                  [
                      [1.0, 0.0],
                      [1.0, 0.0],
                      [1.0, 0.0],
                      [0.0, 0.0],
                      [0.0, 0.0],
                      [0.0, 0.0],
                  ]
              )[None, :, :],
          ],
          dim=0,
      )
      map_states = hmm_distr.map_states(x)
      filter = hmm_distr.filter(x)
      smooth = hmm_distr.smooth(x)
  
      expected_map_states = torch.cat(
          [
              indiv_map_states(a0, a, e, x[0])[None, :],
              indiv_map_states(a0, a, e, x[1])[None, :],
          ],
          -2,
      )
      expected_filter = torch.cat(
          [
              indiv_filter(a0, a, e, x[0])[None, :, :],
              indiv_filter(a0, a, e, x[1])[None, :, :],
          ],
          -3,
      )
      expected_smooth = torch.cat(
          [
              indiv_smooth(a0, a, e, x[0])[None, :, :],
              indiv_smooth(a0, a, e, x[1])[None, :, :],
          ],
          -3,
      )
  
      assert torch.allclose(map_states, expected_map_states)
      assert torch.allclose(filter, expected_filter)
      assert torch.allclose(smooth, expected_smooth)
  
      # Batch parameters.
      a0 = torch.cat(
          [a0[None, :], torch.tensor([0.2, 0.7, 0.1])[None, :]]
      )
      a = torch.cat(
          [
              a[None, :, :],
              torch.tensor(
                  [[0.8, 0.1, 0.1], [0.2, 0.6, 0.2], [0.1, 0.1, 0.8]]
              )[None, :, :],
          ],
          dim=0,
      )
      e = torch.cat(
          [
              e[None, :, :],
              torch.tensor([[0.4, 0.6], [0.99, 0.01], [0.7, 0.3]])[
                  None, :, :
              ],
          ],
          dim=0,
      )
      hmm_distr = MissingDataDiscreteHMM(
          torch.log(a0), torch.log(a), torch.log(e)
      )
      map_states = hmm_distr.map_states(x[1])
      filter = hmm_distr.filter(x[1])
      smooth = hmm_distr.smooth(x[1])
  
      expected_map_states = torch.cat(
          [
              indiv_map_states(a0[0], a[0], e[0], x[1])[None, :],
              indiv_map_states(a0[1], a[1], e[1], x[1])[None, :],
          ],
          -2,
      )
      expected_filter = torch.cat(
          [
              indiv_filter(a0[0], a[0], e[0], x[1])[None, :, :],
              indiv_filter(a0[1], a[1], e[1], x[1])[None, :, :],
          ],
          -3,
      )
      expected_smooth = torch.cat(
          [
              indiv_smooth(a0[0], a[0], e[0], x[1])[None, :, :],
              indiv_smooth(a0[1], a[1], e[1], x[1])[None, :, :],
          ],
          -3,
      )
  
      assert torch.allclose(map_states, expected_map_states)
      assert torch.allclose(filter, expected_filter)
      assert torch.allclose(smooth, expected_smooth)
  
      # Batch both parameters and values.
      map_states = hmm_distr.map_states(x)
      filter = hmm_distr.filter(x)
      smooth = hmm_distr.smooth(x)
  
      expected_map_states = torch.cat(
          [
              indiv_map_states(a0[0], a[0], e[0], x[0])[None, :],
              indiv_map_states(a0[1], a[1], e[1], x[1])[None, :],
          ],
          -2,
      )
      expected_filter = torch.cat(
          [
              indiv_filter(a0[0], a[0], e[0], x[0])[None, :, :],
              indiv_filter(a0[1], a[1], e[1], x[1])[None, :, :],
          ],
          -3,
      )
      expected_smooth = torch.cat(
          [
              indiv_smooth(a0[0], a[0], e[0], x[0])[None, :, :],
              indiv_smooth(a0[1], a[1], e[1], x[1])[None, :, :],
          ],
          -3,
      )
  
      assert torch.allclose(map_states, expected_map_states)
      assert torch.allclose(filter, expected_filter)
      assert torch.allclose(smooth, expected_smooth)
  
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 8 for calling a global function with example usage: test_state_infer() without return types