code: def test_expand_reshaped_distribution(
      extra_event_dims, expand_shape, default
  ):
      probs = torch.ones(1, 6) / 6
      d = dist.OneHotCategorical(probs)
      full_shape = torch.Size([4, 1, 1, 1, 6])
      if default:
          reshaped_dist = TorchDistribution.expand(
              d, [4, 1, 1, 1]
          ).to_event(extra_event_dims)
      else:
          reshaped_dist = d.expand_by([4, 1, 1]).to_event(
              extra_event_dims
          )
      cut = 4 - extra_event_dims
      batch_shape, event_shape = full_shape[:cut], full_shape[cut:]
      assert reshaped_dist.batch_shape == batch_shape
      assert reshaped_dist.event_shape == event_shape
      large = reshaped_dist.expand(expand_shape)
      assert large.batch_shape == torch.Size(expand_shape)
      assert large.event_shape == torch.Size(event_shape)
  
      # Throws error when batch shape cannot be broadcasted
      with pytest.raises((RuntimeError, ValueError)):
          reshaped_dist.expand(expand_shape + [3])
  
      # Throws error when trying to shrink existing batch shape
      with pytest.raises((RuntimeError, ValueError)):
          large.expand(expand_shape[1:])
  
imports: import numpy
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.distributions import TorchDistribution
  from pyro.distributions.testing.gof import auto_goodness_of_fit
  from pyro.distributions.util import broadcast_shape
  from tests.common import TEST_FAILURE_RATE, assert_close, assert_equal, xfail_if_not_implemented
label_desc: Use image node_id 24 for calling a global function with example usage: test_expand_reshaped_distribution(extra_event_dims, expand_shape, default) without return types