code: def test_hmm_example(prev_enum_dim, curr_enum_dim):
      hidden_dim = 8
      probs_x = torch.rand(hidden_dim, hidden_dim, hidden_dim)
      x_prev = torch.arange(hidden_dim).reshape(
          (-1,) + (1,) * (-1 - prev_enum_dim)
      )
      x_curr = torch.arange(hidden_dim).reshape(
          (-1,) + (1,) * (-1 - curr_enum_dim)
      )
  
      expected = probs_x[
          x_prev.unsqueeze(-1),
          x_curr.unsqueeze(-1),
          torch.arange(hidden_dim),
      ]
      actual = Vindex(probs_x)[x_prev, x_curr, :]
      assert_equal(actual, expected)
  
file_path: pyro/tests/ops/test_indexing.py
imports: import itertools
  import pytest
  import torch
  import pyro.distributions
  from pyro.distributions.util import broadcast_shape
  from pyro.ops.indexing import Index, Vindex
  from tests.common import assert_equal
label_desc: Use image node_id 4 for calling a global function with example usage: test_hmm_example(prev_enum_dim, curr_enum_dim) without return types