code: def test_batch_log_dims(dim, probs):
      probs = modify_params_using_dims(probs, dim)
      log_prob_shape = torch.Size(
          (3,) + dist.Categorical(probs).batch_shape
      )
      support = dist.Categorical(probs).enumerate_support()
      log_prob = dist.Categorical(probs).log_prob(support)
      assert_equal(log_prob.size(), log_prob_shape)
  
imports: from unittest import TestCase
  import numpy
  import pytest
  import scipy.stats
  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_batch_log_dims(dim, probs) without return types