code: def test_support_dims(dim, probs):
      probs = modify_params_using_dims(probs, dim)
      d = dist.OneHotCategorical(probs)
      support = d.enumerate_support()
      for s in support:
          assert_correct_dimensions(s, probs)
      n = len(support)
      assert support.shape == (n,) + d.batch_shape + d.event_shape
      support_expanded = d.enumerate_support(expand=True)
      assert (
          support_expanded.shape == (n,) + d.batch_shape + d.event_shape
      )
      support_unexpanded = d.enumerate_support(expand=False)
      assert (
          support_unexpanded.shape
          == (n,) + (1,) * len(d.batch_shape) + d.event_shape
      )
  
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 6 for calling a global function with example usage: test_support_dims(dim, probs) without return types