code: def test_batch_log_dims(dim, probs):
      batch_pdf_shape = (3,) + (1,) * (dim - 1)
      expected_log_prob_sum = np.array(
          wrap_nested(list(np.log(probs)), dim - 1)
      ).reshape(*batch_pdf_shape)
      probs = modify_params_using_dims(probs, dim)
      support = dist.OneHotCategorical(probs).enumerate_support()
      log_prob = dist.OneHotCategorical(probs).log_prob(support)
      assert_equal(
          log_prob.detach().cpu().numpy(), expected_log_prob_sum
      )
  
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 8 for calling a global function with example usage: test_batch_log_dims(dim, probs) without return types