code: def test_one_hot_categorical_enumerate(shape, expand):
      shape = torch.Size(shape)
      probs = torch.ones(shape)
  
      @pyro.ops.jit.trace
      def f(probs):
          d = dist.OneHotCategorical(probs)
          support = d.enumerate_support(expand=expand)
          return d.log_prob(support)
  
      log_prob = f(probs)
      batch_shape = shape[:-1]
      assert log_prob.shape == shape[-1:] + batch_shape
  
imports: import logging
  import warnings
  import pytest
  import torch
  from torch.autograd import grad
  from torch.distributions import constraints, kl_divergence
  import pyro
  import pyro.distributions
  import pyro.ops.jit
  import pyro.poutine
  from pyro.distributions.util import scale_and_mask
  from pyro.infer import SVI, JitTrace_ELBO, JitTraceEnum_ELBO, JitTraceGraph_ELBO, JitTraceMeanField_ELBO, Trace_ELBO, TraceEnum_ELBO, TraceGraph_ELBO, TraceMeanField_ELBO, TraceTMC_ELBO, infer_discrete
  from pyro.optim import Adam
  from pyro.poutine.indep_messenger import CondIndepStackFrame
  from pyro.util import ignore_jit_warnings
  from tests.common import assert_close, assert_equal
label_desc: Use image node_id 13 for calling a global function with example usage: test_one_hot_categorical_enumerate(shape, expand) and returns: d