code: def test_sample_dims(dim, probs):
      probs = modify_params_using_dims(probs, dim)
      sample = dist.OneHotCategorical(probs).sample()
      assert_correct_dimensions(sample, probs)
  
file_path: pyro/tests/distributions/test_one_hot_categorical.py
imports: from unittest import TestCase
  import numpy
  import pytest
  import torch
  import pyro.distributions
  from tests.common import assert_equal
label_desc: Use image node_id 7 for calling a global function with example usage: test_sample_dims(dim, probs) without return types