code: def test_reshape(sample_dim, extra_event_dims):
      batch_dim = 3
      batch_shape, event_shape = torch.Size((5, 4, 3)), torch.Size()
      sample_shape = torch.Size((8, 7, 6))[3 - sample_dim :]
      shape = sample_shape + batch_shape + event_shape
  
      # Construct a base dist of desired starting shape.
      dist0 = Bernoulli(0.5 * torch.ones(batch_shape))
      assert dist0.event_shape == event_shape
      assert dist0.batch_shape == batch_shape
  
      # Check that reshaping has the desired final shape.
      dist = dist0.expand_by(sample_shape).to_event(extra_event_dims)
      sample = dist.sample()
      assert sample.shape == shape
      assert dist.mean.shape == shape
      assert dist.variance.shape == shape
      assert (
          dist.log_prob(sample).shape
          == shape[: sample_dim + batch_dim - extra_event_dims]
      )
  
      # Check enumerate support.
      if dist.event_shape:
          with pytest.raises(NotImplementedError):
              dist.enumerate_support()
          with pytest.raises(NotImplementedError):
              dist.enumerate_support(expand=True)
          with pytest.raises(NotImplementedError):
              dist.enumerate_support(expand=False)
      else:
          assert dist.enumerate_support().shape == (2,) + shape
          assert (
              dist.enumerate_support(expand=True).shape == (2,) + shape
          )
          assert (
              dist.enumerate_support(expand=False).shape
              == (2,)
              + (1,) * len(sample_shape + batch_shape)
              + event_shape
          )
  
file_path: pyro/tests/distributions/test_reshape.py
imports: import pytest
  import torch
  from pyro.distributions.torch import Bernoulli
  from tests.common import assert_equal
label_desc: Use image node_id 3 for calling a global function with example usage: test_reshape(sample_dim, extra_event_dims) without return types